Spaces:
Sleeping
Sleeping
File size: 4,232 Bytes
4dbe519 64462d2 4dbe519 378dbdf 4dbe519 378dbdf 4dbe519 378dbdf 4dbe519 64462d2 4dbe519 378dbdf 4dbe519 378dbdf 4dbe519 378dbdf 4dbe519 378dbdf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """LLM client and request/response models."""
import json
from typing import Any, List, Optional, Dict
from pydantic import BaseModel, Field, ConfigDict
from litellm import acompletion
from .models import Message, ToolCall, ToolResult, ContentItem
def build_messages(request: 'LlmRequest') -> List[dict]:
"""Convert LlmRequest to API message format.
Standalone function for use by memory/callback modules.
"""
messages = []
for instruction in request.instructions:
messages.append({"role": "system", "content": instruction})
for item in request.contents:
if isinstance(item, Message):
messages.append({"role": item.role, "content": item.content})
elif isinstance(item, ToolCall):
tool_call_dict = {
"id": item.tool_call_id,
"type": "function",
"function": {
"name": item.name,
"arguments": json.dumps(item.arguments)
}
}
# Append to previous assistant message if exists
if messages and messages[-1]["role"] == "assistant":
messages[-1].setdefault("tool_calls", []).append(tool_call_dict)
else:
messages.append({
"role": "assistant",
"content": None,
"tool_calls": [tool_call_dict]
})
elif isinstance(item, ToolResult):
messages.append({
"role": "tool",
"tool_call_id": item.tool_call_id,
"content": str(item.content[0]) if item.content else ""
})
return messages
class LlmRequest(BaseModel):
"""Request object for LLM calls."""
instructions: List[str] = Field(default_factory=list)
contents: List[ContentItem] = Field(default_factory=list)
tools: List[Any] = Field(default_factory=list)
tool_choice: Optional[str] = 'auto'
class LlmResponse(BaseModel):
"""Response object from LLM calls."""
content: List[ContentItem] = Field(default_factory=list)
error_message: Optional[str] = None
usage_metadata: Dict[str, Any] = Field(default_factory=dict)
class LlmClient:
"""Client for LLM API calls using LiteLLM."""
def __init__(self, model: str, **config):
self.model = model
self.config = config
async def generate(self, request: LlmRequest) -> LlmResponse:
"""Generate a response from the LLM."""
try:
messages = self._build_messages(request)
tools = [t.tool_definition for t in request.tools] if request.tools else None
response = await acompletion(
model=self.model,
messages=messages,
tools=tools,
**({"tool_choice": request.tool_choice}
if request.tool_choice else {}),
**self.config
)
return self._parse_response(response)
except Exception as e:
return LlmResponse(error_message=str(e))
def _build_messages(self, request: LlmRequest) -> List[dict]:
"""Convert LlmRequest to API message format."""
return build_messages(request)
def _parse_response(self, response) -> LlmResponse:
"""Convert API response to LlmResponse."""
choice = response.choices[0]
content_items = []
if choice.message.content:
content_items.append(Message(
role="assistant",
content=choice.message.content
))
if choice.message.tool_calls:
for tc in choice.message.tool_calls:
content_items.append(ToolCall(
tool_call_id=tc.id,
name=tc.function.name,
arguments=json.loads(tc.function.arguments)
))
return LlmResponse(
content=content_items,
usage_metadata={
"input_tokens": response.usage.prompt_tokens,
"output_tokens": response.usage.completion_tokens,
}
) |