wehub-resource-sync
87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
"""Llamaindex LLM implementation of Claude Vertex AI"""
|
|
|
|
from typing import Any
|
|
|
|
from anthropic import AnthropicVertex, AsyncAnthropicVertex
|
|
from llama_index.core.llms import (
|
|
CompletionResponse,
|
|
CompletionResponseGen,
|
|
CustomLLM,
|
|
LLMMetadata,
|
|
)
|
|
from llama_index.core.llms.callbacks import llm_completion_callback
|
|
from pydantic import Field, PrivateAttr
|
|
|
|
|
|
class ClaudeVertexLLM(CustomLLM):
|
|
project_id: str = Field(description="The project ID for Vertex AI")
|
|
region: str = Field(description="The region for Vertex AI")
|
|
model_name: str = Field(description="The name of the Claude model to use")
|
|
max_tokens: int = Field(
|
|
description="The maximum number \
|
|
of tokens to generate"
|
|
)
|
|
system_prompt: str = Field(description="The system prompt to use")
|
|
|
|
client: Any = PrivateAttr()
|
|
async_client: Any = PrivateAttr()
|
|
|
|
def __init__(self, **data):
|
|
super().__init__(**data)
|
|
self.client = AnthropicVertex(project_id=self.project_id, region=self.region)
|
|
self.async_client = AsyncAnthropicVertex(
|
|
project_id=self.project_id, region=self.region
|
|
)
|
|
|
|
@property
|
|
def metadata(self) -> LLMMetadata:
|
|
"""Get LLM metadata."""
|
|
return LLMMetadata(
|
|
model_name=self.model_name,
|
|
max_tokens=self.max_tokens,
|
|
system_prompt=self.system_prompt,
|
|
)
|
|
|
|
@llm_completion_callback()
|
|
def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
|
|
message = self.client.messages.create(
|
|
model=self.model_name,
|
|
max_tokens=self.max_tokens,
|
|
system=self.system_prompt,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": prompt,
|
|
}
|
|
],
|
|
)
|
|
return CompletionResponse(text=message.content[0].text)
|
|
|
|
@llm_completion_callback()
|
|
async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
|
|
message = await self.async_client.messages.create(
|
|
model=self.model_name,
|
|
max_tokens=self.max_tokens,
|
|
system=self.system_prompt,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": prompt,
|
|
}
|
|
],
|
|
)
|
|
return CompletionResponse(text=message.content[0].text)
|
|
|
|
@llm_completion_callback()
|
|
def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
|
|
with self.client.messages.stream(
|
|
model=self.model_name,
|
|
max_tokens=self.max_tokens,
|
|
system=self.system_prompt,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
) as stream:
|
|
response = ""
|
|
for text in stream.text_stream:
|
|
response += text
|
|
yield CompletionResponse(text=response, delta=text)
|