38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
from typing import AsyncGenerator
|
|
from anthropic import AsyncAnthropic
|
|
from dailyai.pipeline.frames import Frame, LLMMessagesFrame, TextFrame
|
|
|
|
from dailyai.services.ai_services import LLMService
|
|
|
|
|
|
class AnthropicLLMService(LLMService):
|
|
|
|
def __init__(
|
|
self,
|
|
api_key,
|
|
model="claude-3-opus-20240229",
|
|
max_tokens=1024):
|
|
super().__init__()
|
|
self.client = AsyncAnthropic(api_key=api_key)
|
|
self.model = model
|
|
self.max_tokens = max_tokens
|
|
|
|
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
|
|
if not isinstance(frame, LLMMessagesFrame):
|
|
yield frame
|
|
|
|
stream = await self.client.messages.create(
|
|
max_tokens=self.max_tokens,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": "Hello, Claude",
|
|
}
|
|
],
|
|
model=self.model,
|
|
stream=True,
|
|
)
|
|
async for event in stream:
|
|
if event.type == "content_block_delta":
|
|
yield TextFrame(event.delta.text)
|