cai/examples/model_providers/litellm.py

56 lines
1.6 KiB
Python

import os
from dotenv import load_dotenv
from openai import AsyncOpenAI
from cai.sdk.agents import OpenAIChatCompletionsModel, Agent, Runner
from cai.sdk.agents import set_default_openai_client, set_tracing_disabled
from openai.types.responses import ResponseTextDeltaEvent
# Load environment variables from .env file
load_dotenv()
external_client = AsyncOpenAI(
base_url=os.getenv("LITELLM_BASE_URL", "http://localhost:4000"),
api_key=os.getenv("LITELLM_API_KEY", "key"),
)
set_default_openai_client(external_client)
set_tracing_disabled(True)
# llm_model=os.getenv('LLM_MODEL', 'gpt-4o')
# llm_model=os.getenv('LLM_MODEL', 'claude-3-7')
llm_model = os.getenv("LLM_MODEL", "qwen2.5:14b")
# For Qwen models, we need to skip system instructions as they're not supported
instructions = None if "qwen" in llm_model.lower() else "You are a helpful assistant"
agent = Agent(
name="Assistant",
instructions=instructions,
model=OpenAIChatCompletionsModel(
model=llm_model,
openai_client=external_client,
),
)
def run_sync_haiku():
result = Runner.run_sync(agent, "Write a haiku about recursion in programming.")
print(result.final_output)
async def stream_jokes():
result = Runner.run_streamed(agent, input="Please tell me 5 jokes.")
async for event in result.stream_events():
if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
print(event.data.delta, end="", flush=True)
########################################################
# async
# import asyncio
# asyncio.run(stream_jokes())
# sync
run_sync_haiku()