lite-agent 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

lite_agent/stream_handlers/__init__.py

@@ -1,5 +1,6 @@
-from lite_agent.stream_handlers.litellm import litellm_stream_handler
+from lite_agent.stream_handlers.litellm import litellm_completion_stream_handler, litellm_response_stream_handler
 
 __all__ = [
-    "litellm_stream_handler",
+    "litellm_completion_stream_handler",
+    "litellm_response_stream_handler",
 ]

lite_agent/stream_handlers/litellm.py

@@ -1,14 +1,19 @@
 from collections.abc import AsyncGenerator
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import aiofiles
 import litellm
-from aiofiles.threadpool.text import AsyncTextIOWrapper
-from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+from litellm.types.llms.openai import ResponsesAPIStreamingResponse
+from litellm.types.utils import ModelResponseStream
+from pydantic import BaseModel
 
 from lite_agent.loggers import logger
-from lite_agent.processors import StreamChunkProcessor
-from lite_agent.types import AgentChunk, CompletionRawChunk, ContentDeltaChunk, FinalMessageChunk, ToolCallDeltaChunk, UsageChunk
+from lite_agent.processors import CompletionEventProcessor, ResponseEventProcessor
+from lite_agent.types import AgentChunk
+
+if TYPE_CHECKING:
+    from aiofiles.threadpool.text import AsyncTextIOWrapper
 
 
 def ensure_record_file(record_to: Path | None) -> Path | None:
@@ -20,75 +25,14 @@ def ensure_record_file(record_to: Path | None) -> Path | None:
     return record_to
 
 
-async def process_chunk(
-    processor: StreamChunkProcessor,
-    chunk: ModelResponseStream,
-    record_file: AsyncTextIOWrapper | None = None,
-) -> AsyncGenerator[AgentChunk, None]:
-    if record_file:
-        await record_file.write(chunk.model_dump_json() + "\n")
-        await record_file.flush()
-    yield CompletionRawChunk(type="completion_raw", raw=chunk)
-    usage_chunk = await handle_usage_chunk(processor, chunk)
-    if usage_chunk:
-        yield usage_chunk
-        return
-    if not chunk.choices:
-        return
-    choice = chunk.choices[0]
-    delta = choice.delta
-    for result in await handle_content_and_tool_calls(processor, chunk, choice, delta):
-        yield result
-    if choice.finish_reason:
-        current_message = processor.current_message
-        yield FinalMessageChunk(type="final_message", message=current_message, finish_reason=choice.finish_reason)
-
-
-async def handle_usage_chunk(processor: StreamChunkProcessor, chunk: ModelResponseStream) -> UsageChunk | None:
-    usage = processor.handle_usage_info(chunk)
-    if usage:
-        return UsageChunk(type="usage", usage=usage)
-    return None
-
-
-async def handle_content_and_tool_calls(
-    processor: StreamChunkProcessor,
-    chunk: ModelResponseStream,
-    choice: StreamingChoices,
-    delta: Delta,
-) -> list[AgentChunk]:
-    results: list[AgentChunk] = []
-    if not processor.is_initialized:
-        processor.initialize_message(chunk, choice)
-    if delta.content:
-        results.append(ContentDeltaChunk(type="content_delta", delta=delta.content))
-        processor.update_content(delta.content)
-    if delta.tool_calls is not None:
-        processor.update_tool_calls(delta.tool_calls)
-        if delta.tool_calls and processor.current_message.tool_calls:
-            results.extend(
-                [
-                    ToolCallDeltaChunk(
-                        type="tool_call_delta",
-                        tool_call_id=processor.current_message.tool_calls[-1].id,
-                        name=processor.current_message.tool_calls[-1].function.name,
-                        arguments_delta=tool_call.function.arguments or "",
-                    )
-                    for tool_call in delta.tool_calls
-                    if tool_call.function.arguments
-                ],
-            )
-    return results
-
-
-async def litellm_stream_handler(
+async def litellm_completion_stream_handler(
     resp: litellm.CustomStreamWrapper,
     record_to: Path | None = None,
 ) -> AsyncGenerator[AgentChunk, None]:
     """
     Optimized chunk handler
     """
-    processor = StreamChunkProcessor()
+    processor = CompletionEventProcessor()
     record_file: AsyncTextIOWrapper | None = None
     record_path = ensure_record_file(record_to)
     if record_path:
@@ -99,7 +43,32 @@ async def litellm_stream_handler(
                 logger.warning("unexpected chunk type: %s", type(chunk))
                 logger.warning("chunk content: %s", chunk)
                 continue
-            async for result in process_chunk(processor, chunk, record_file):
+            async for result in processor.process_chunk(chunk, record_file):
+                yield result
+    finally:
+        if record_file:
+            await record_file.close()
+
+
+async def litellm_response_stream_handler(
+    resp: AsyncGenerator[ResponsesAPIStreamingResponse, None],
+    record_to: Path | None = None,
+) -> AsyncGenerator[AgentChunk, None]:
+    """
+    Response API stream handler for processing ResponsesAPIStreamingResponse chunks
+    """
+    processor = ResponseEventProcessor()
+    record_file: AsyncTextIOWrapper | None = None
+    record_path = ensure_record_file(record_to)
+    if record_path:
+        record_file = await aiofiles.open(record_path, "w", encoding="utf-8")
+    try:
+        async for chunk in resp:
+            if not isinstance(chunk, BaseModel):
+                logger.warning("unexpected chunk type: %s", type(chunk))
+                logger.warning("chunk content: %s", chunk)
+                continue
+            async for result in processor.process_chunk(chunk, record_file):
                 yield result
     finally:
         if record_file:

lite_agent/types/__init__.py

@@ -1,39 +1,70 @@
 # Export all types from submodules
-from .chunks import (
+from .events import (
     AgentChunk,
     AgentChunkType,
-    CompletionRawChunk,
-    ContentDeltaChunk,
-    FinalMessageChunk,
-    ToolCallChunk,
-    ToolCallDeltaChunk,
-    ToolCallResultChunk,
-    UsageChunk,
+    AssistantMessageEvent,
+    CompletionRawEvent,
+    ContentDeltaEvent,
+    FunctionCallDeltaEvent,
+    FunctionCallEvent,
+    FunctionCallOutputEvent,
+    ResponseRawEvent,
+    Timing,
+    TimingEvent,
+    UsageEvent,
+)
+from .events import (
+    Usage as EventUsage,
 )
 from .messages import (
+    # Legacy compatibility aliases (temporary)
     AgentAssistantMessage,
-    AgentFunctionCallOutput,
-    AgentFunctionToolCallMessage,
-    AgentMessage,
     AgentSystemMessage,
     AgentUserMessage,
     AssistantMessage,
+    # New types
+    AssistantMessageContent,
     AssistantMessageDict,
+    AssistantMessageMeta,
+    AssistantTextContent,
+    AssistantToolCall,
+    AssistantToolCallResult,
+    BasicMessageMeta,
     FlexibleRunnerMessage,
     FunctionCallDict,
     FunctionCallOutputDict,
-    Message,
+    LLMResponseMeta,
     MessageDict,
+    # New metadata types
+    MessageMeta,
+    MessageUsage,
+    NewAssistantMessage,
+    NewMessage,
+    NewMessages,
+    NewSystemMessage,
+    # New structured message types
+    NewUserMessage,
     ResponseInputImage,
     ResponseInputText,
     RunnerMessage,
     RunnerMessages,
     SystemMessageDict,
+    UserFileContent,
+    UserImageContent,
     UserInput,
+    # New content types
+    UserMessageContent,
     UserMessageContentItemImageURL,
     UserMessageContentItemImageURLImageURL,
     UserMessageContentItemText,
     UserMessageDict,
+    UserTextContent,
+    # Utility functions
+    assistant_message_to_llm_dict,
+    message_to_llm_dict,
+    messages_to_llm_format,
+    system_message_to_llm_dict,
+    user_message_to_llm_dict,
 )
 from .tool_calls import ToolCall, ToolCallFunction
 
@@ -41,35 +72,58 @@ __all__ = [
     "AgentAssistantMessage",
     "AgentChunk",
     "AgentChunkType",
-    "AgentFunctionCallOutput",
-    "AgentFunctionToolCallMessage",
-    "AgentMessage",
     "AgentSystemMessage",
     "AgentUserMessage",
     "AssistantMessage",
+    "AssistantMessageContent",
     "AssistantMessageDict",
-    "CompletionRawChunk",
-    "ContentDeltaChunk",
-    "FinalMessageChunk",
+    "AssistantMessageEvent",
+    "AssistantMessageMeta",
+    "AssistantTextContent",
+    "AssistantToolCall",
+    "AssistantToolCallResult",
+    "BasicMessageMeta",
+    "CompletionRawEvent",
+    "ContentDeltaEvent",
+    "EventUsage",
     "FlexibleRunnerMessage",
+    "FunctionCallDeltaEvent",
     "FunctionCallDict",
+    "FunctionCallEvent",
     "FunctionCallOutputDict",
-    "Message",
+    "FunctionCallOutputEvent",
+    "LLMResponseMeta",
     "MessageDict",
+    "MessageMeta",
+    "MessageUsage",
+    "NewAssistantMessage",
+    "NewMessage",
+    "NewMessages",
+    "NewSystemMessage",
+    "NewUserMessage",
     "ResponseInputImage",
     "ResponseInputText",
+    "ResponseRawEvent",
     "RunnerMessage",
     "RunnerMessages",
     "SystemMessageDict",
+    "Timing",
+    "TimingEvent",
     "ToolCall",
-    "ToolCallChunk",
-    "ToolCallDeltaChunk",
     "ToolCallFunction",
-    "ToolCallResultChunk",
-    "UsageChunk",
+    "UsageEvent",
+    "UserFileContent",
+    "UserImageContent",
     "UserInput",
+    "UserMessageContent",
     "UserMessageContentItemImageURL",
     "UserMessageContentItemImageURLImageURL",
     "UserMessageContentItemText",
     "UserMessageDict",
+    "UserTextContent",
+    "assistant_message_to_llm_dict",
+    "message_to_llm_dict",
+    "messages_to_llm_format",
+    "system_message_to_llm_dict",
+    "user_message_to_llm_dict",
 ]

lite_agent/types/events.py

@@ -0,0 +1,119 @@
+from typing import Literal
+
+from litellm.types.utils import ModelResponseStream
+from pydantic import BaseModel
+
+from .messages import NewAssistantMessage
+
+
+class Usage(BaseModel):
+    input_tokens: int
+    output_tokens: int
+
+
+class Timing(BaseModel):
+    latency_ms: int
+    output_time_ms: int
+
+
+class CompletionRawEvent(BaseModel):
+    """
+    Define the type of chunk
+    """
+
+    type: Literal["completion_raw"] = "completion_raw"
+    raw: ModelResponseStream
+
+
+class ResponseRawEvent(BaseModel):
+    """
+    Define the type of response raw chunk
+    """
+
+    type: Literal["response_raw"] = "response_raw"
+    raw: object
+
+
+class UsageEvent(BaseModel):
+    """
+    Define the type of usage info chunk
+    """
+
+    type: Literal["usage"] = "usage"
+    usage: Usage
+
+
+class TimingEvent(BaseModel):
+    """
+    Define the type of timing info chunk
+    """
+
+    type: Literal["timing"] = "timing"
+    timing: Timing
+
+
+class AssistantMessageEvent(BaseModel):
+    """
+    Define the type of assistant message chunk
+    """
+
+    type: Literal["assistant_message"] = "assistant_message"
+    message: NewAssistantMessage
+
+
+class FunctionCallEvent(BaseModel):
+    """
+    Define the type of tool call chunk
+    """
+
+    type: Literal["function_call"] = "function_call"
+    call_id: str
+    name: str
+    arguments: str
+
+
+class FunctionCallOutputEvent(BaseModel):
+    """
+    Define the type of tool call result chunk
+    """
+
+    type: Literal["function_call_output"] = "function_call_output"
+    tool_call_id: str
+    name: str
+    content: str
+    execution_time_ms: int | None = None
+
+
+class ContentDeltaEvent(BaseModel):
+    """
+    Define the type of message chunk
+    """
+
+    type: Literal["content_delta"] = "content_delta"
+    delta: str
+
+
+class FunctionCallDeltaEvent(BaseModel):
+    """
+    Define the type of tool call delta chunk
+    """
+
+    type: Literal["function_call_delta"] = "function_call_delta"
+    tool_call_id: str
+    name: str
+    arguments_delta: str
+
+
+AgentChunk = CompletionRawEvent | ResponseRawEvent | UsageEvent | TimingEvent | FunctionCallEvent | FunctionCallOutputEvent | ContentDeltaEvent | FunctionCallDeltaEvent | AssistantMessageEvent
+
+AgentChunkType = Literal[
+    "completion_raw",
+    "response_raw",
+    "usage",
+    "timing",
+    "function_call",
+    "function_call_output",
+    "content_delta",
+    "function_call_delta",
+    "assistant_message",
+]