lite-agent 0.1.0__py3-none-any.whl

lite_agent/__main__.py

@@ -0,0 +1,110 @@
+import asyncio
+import logging
+
+from prompt_toolkit import PromptSession
+from prompt_toolkit.validation import Validator
+from rich.console import Console
+from rich.logging import RichHandler
+
+from open_agents.agent import Agent
+from open_agents.loggers import logger
+from open_agents.runner import Runner
+from open_agents.types import AgentChunk, ContentDeltaChunk
+
+logging.basicConfig(level=logging.WARNING, handlers=[RichHandler()], format="%(message)s")
+logger.setLevel(logging.INFO)
+
+
+# --- Tool functions (English) ---
+async def get_whether(city: str) -> str:
+    """Get the weather for a city."""
+    await asyncio.sleep(1)  # 模擬網路延遲
+    return f"The weather in {city} is sunny with a few clouds."
+
+
+async def get_temperature(city: str) -> str:
+    """Get the temperature for a city."""
+    await asyncio.sleep(1)  # 模擬網路延遲
+    return f"The temperature in {city} is 25°C."
+
+
+class RichChannel:
+    def __init__(self) -> None:
+        self.console = Console()
+        self.map = {
+            "final_message": self.handle_final_message,
+            "tool_call": self.handle_tool_call,
+            "tool_call_result": self.handle_tool_call_result,
+            "tool_call_delta": self.handle_tool_call_delta,
+            "content_delta": self.handle_content_delta,
+            "usage": self.handle_usage,
+        }
+        self.new_turn = True
+
+    def handle(self, chunk: AgentChunk):
+        handler = self.map[chunk["type"]]
+        handler(chunk)
+
+    def handle_final_message(self, _chunk: AgentChunk):
+        print()
+        self.new_turn = True
+
+    def handle_tool_call(self, chunk: AgentChunk):
+        name = chunk.get("name", "<unknown>")
+        arguments = chunk.get("arguments", "")
+        self.console.print(f"🛠️  [green]{name}[/green]([yellow]{arguments}[/yellow])")
+
+    def handle_tool_call_result(self, chunk: AgentChunk):
+        name = chunk.get("name", "<unknown>")
+        content = chunk.get("content", "")
+        self.console.print(f"🛠️  [green]{name}[/green] → [yellow]{content}[/yellow]")
+
+    def handle_tool_call_delta(self, chunk: AgentChunk): ...
+    def handle_content_delta(self, chunk: ContentDeltaChunk):
+        if self.new_turn:
+            self.console.print("🤖 ", end="")
+            self.new_turn = False
+        print(chunk["delta"], end="", flush=True)
+
+    def handle_usage(self, chunk: AgentChunk):
+        if False:
+            usage = chunk["usage"]
+            self.console.print(f"In: {usage.prompt_tokens}, Out: {usage.completion_tokens}, Total: {usage.total_tokens}")
+
+
+async def main():
+    agent = Agent(
+        model="gpt-4.1",
+        name="Weather Assistant",
+        instructions="You are a helpful weather assistant. Before using tools, briefly explain what you are going to do. Provide friendly and informative responses.",
+        tools=[get_whether, get_temperature],
+    )
+    session = PromptSession()
+    rich_channel = RichChannel()
+    runner = Runner(agent)
+    not_empty_validator = Validator.from_callable(
+        lambda text: bool(text.strip()),
+        error_message="Input cannot be empty.",
+        move_cursor_to_end=True,
+    )
+    while True:
+        try:
+            user_input = await session.prompt_async(
+                "👤 ",
+                default="",
+                complete_while_typing=True,
+                validator=not_empty_validator,
+                validate_while_typing=False,
+            )
+            if user_input.lower() in {"exit", "quit"}:
+                break
+            response = runner.run_stream(user_input)
+            async for chunk in response:
+                rich_channel.handle(chunk)
+
+        except (EOFError, KeyboardInterrupt):
+            break
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

lite_agent/agent.py

@@ -0,0 +1,36 @@
+from collections.abc import AsyncGenerator, Callable
+
+import litellm
+from funcall import Funcall
+
+from open_agents.chunk_handler import AgentChunk, chunk_handler
+from open_agents.types import RunnerMessages
+
+
+class Agent:
+    def __init__(self, *, model: str, name: str, instructions: str, tools: list[Callable] | None = None) -> None:
+        self.name = name
+        self.instructions = instructions
+        self.fc = Funcall(tools)
+        self.model = model
+
+    def prepare_messages(self, messages: RunnerMessages) -> list[dict]:
+        return [
+            {
+                "role": "system",
+                "content": f"You are {self.name}. {self.instructions}",
+            },
+            *messages,
+        ]
+
+    async def stream_async(self, messages: RunnerMessages) -> AsyncGenerator[AgentChunk, None]:
+        self.message_histories = self.prepare_messages(messages)
+        tools = self.fc.get_tools(target="litellm")
+        resp = await litellm.acompletion(
+            model=self.model,
+            messages=self.message_histories,
+            tools=tools,
+            tool_choice="auto",
+            stream=True,
+        )
+        return chunk_handler(resp, self.fc)

lite_agent/chunk_handler.py

@@ -0,0 +1,166 @@
+from collections.abc import AsyncGenerator
+from typing import Literal, TypedDict
+
+import litellm
+from funcall import Funcall
+
+from open_agents.loggers import logger
+from open_agents.processors import StreamChunkProcessor
+from open_agents.processors.stream_chunk_processor import AssistantMessage
+
+
+class LiteLLMRawChunk(TypedDict):
+    """
+    Define the type of chunk
+    """
+
+    type: Literal["litellm_raw"]
+    raw: litellm.ModelResponseStream
+
+
+class UsageChunk(TypedDict):
+    """
+    Define the type of usage info chunk
+    """
+
+    type: Literal["usage"]
+    usage: litellm.Usage
+
+
+class FinalMessageChunk(TypedDict):
+    """
+    Define the type of final message chunk
+    """
+
+    type: Literal["final_message"]
+    message: AssistantMessage
+    finish_reason: Literal["stop", "tool_calls"]
+
+
+class ToolCallChunk(TypedDict):
+    """
+    Define the type of tool call chunk
+    """
+
+    type: Literal["tool_call"]
+    name: str
+    arguments: str
+
+
+class ToolCallResultChunk(TypedDict):
+    """
+    Define the type of tool call result chunk
+    """
+
+    type: Literal["tool_call_result"]
+    tool_call_id: str
+    name: str
+    content: str
+
+
+class ContentDeltaChunk(TypedDict):
+    """
+    Define the type of message chunk
+    """
+
+    type: Literal["content_delta"]
+    delta: str
+
+
+class ToolCallDeltaChunk(TypedDict):
+    """
+    Define the type of tool call delta chunk
+    """
+
+    type: Literal["tool_call_delta"]
+    tool_call_id: str
+    name: str
+    arguments_delta: str
+
+
+AgentChunk = LiteLLMRawChunk | UsageChunk | FinalMessageChunk | ToolCallChunk | ToolCallResultChunk | ContentDeltaChunk
+
+
+async def chunk_handler(
+    resp: litellm.CustomStreamWrapper,
+    fc: Funcall,
+) -> AsyncGenerator[AgentChunk, None]:
+    """
+    Optimized chunk handler
+
+    Args:
+        resp: LiteLLM streaming response wrapper
+        fc: function call handler
+
+    Yields:
+        litellm.ModelResponseStream: processed response chunk
+
+    Raises:
+        Exception: various exceptions during processing
+    """
+    processor = StreamChunkProcessor(fc)
+    async for chunk in resp:
+        if not isinstance(chunk, litellm.ModelResponseStream):
+            logger.debug("unexpected chunk type: %s", type(chunk))
+            logger.debug("chunk content: %s", chunk)
+            continue
+
+        # Handle usage info
+        if usage := processor.handle_usage_info(chunk):
+            yield UsageChunk(type="usage", usage=usage)
+            continue
+
+        # Get choice and delta data
+        if not chunk.choices:
+            yield LiteLLMRawChunk(type="litellm_raw", raw=chunk)
+            continue
+
+        choice = chunk.choices[0]
+        delta = choice.delta
+        if not processor.current_message:
+            processor.initialize_message(chunk, choice)
+        if delta.content:
+            yield ContentDeltaChunk(type="content_delta", delta=delta.content)
+        processor.update_content(delta.content)
+        processor.update_tool_calls(delta.tool_calls)
+        if delta.tool_calls:
+            for tool_call in delta.tool_calls:
+                if tool_call.function.arguments:
+                    yield 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,
+                    )
+        # Check if finished
+        if choice.finish_reason and processor.current_message:
+            current_message = processor.finalize_message()
+            yield FinalMessageChunk(type="final_message", message=current_message, finish_reason=choice.finish_reason)
+            # New: check tool_calls and handle
+            tool_calls = current_message.tool_calls
+            if tool_calls:
+                # Execute each tool_call and yield result
+                for tool_call in tool_calls:
+                    try:
+                        yield ToolCallChunk(
+                            type="tool_call",
+                            name=tool_call.function.name,
+                            arguments=tool_call.function.arguments,
+                        )
+                        content = await fc.call_function_async(tool_call.function.name, tool_call.function.arguments)
+                        yield ToolCallResultChunk(
+                            type="tool_call_result",
+                            tool_call_id=tool_call.id,
+                            name=tool_call.function.name,
+                            content=str(content),
+                        )
+                    except Exception as e:  # noqa: PERF203
+                        logger.exception("Tool call %s failed", tool_call.id)
+                        yield ToolCallResultChunk(
+                            type="tool_call_result",
+                            tool_call_id=tool_call.id,
+                            name=tool_call.function.name,
+                            content=str(e),
+                        )
+            continue
+        yield LiteLLMRawChunk(type="litellm_raw", raw=chunk)

lite_agent/loggers.py

@@ -0,0 +1,3 @@
+import logging
+
+logger = logging.getLogger("easy_agent")

lite_agent/processors/__init__.py

@@ -0,0 +1,3 @@
+from open_agents.processors.stream_chunk_processor import StreamChunkProcessor
+
+__all__ = ["StreamChunkProcessor"]

lite_agent/processors/stream_chunk_processor.py

@@ -0,0 +1,85 @@
+import litellm
+from funcall import Funcall
+from litellm.types.utils import ChatCompletionDeltaToolCall, StreamingChoices
+
+from open_agents.loggers import logger
+from open_agents.types import AssistantMessage, ToolCall, ToolCallFunction
+
+
+class StreamChunkProcessor:
+    """Processor for handling streaming responses"""
+
+    def __init__(self, fc: Funcall) -> None:
+        self.fc = fc
+        self.current_message: AssistantMessage = None
+
+    def initialize_message(self, chunk: litellm.ModelResponseStream, choice: StreamingChoices) -> None:
+        """Initialize the message object"""
+        delta = choice.delta
+        self.current_message = AssistantMessage(
+            id=chunk.id,
+            index=choice.index,
+            role=delta.role,
+            content="",
+        )
+        logger.debug("Initialized new message: %s", self.current_message.id)
+
+    def update_content(self, content: str) -> None:
+        """Update message content"""
+        if self.current_message and content:
+            self.current_message.content += content
+
+    def _initialize_tool_calls(self, tool_calls: list[litellm.ChatCompletionMessageToolCall]) -> None:
+        """Initialize tool calls"""
+        if not self.current_message:
+            return
+
+        self.current_message.tool_calls = []
+        for call in tool_calls:
+            logger.debug("Create new tool call: %s", call.id)
+
+    def _update_tool_calls(self, tool_calls: list[litellm.ChatCompletionMessageToolCall]) -> None:
+        """Update existing tool calls"""
+        if not self.current_message or not self.current_message.tool_calls:
+            return
+
+        for current_call, new_call in zip(self.current_message.tool_calls, tool_calls, strict=False):
+            if new_call.function.arguments:
+                current_call.function.arguments += new_call.function.arguments
+            if new_call.type:
+                current_call.type = new_call.type
+
+    def update_tool_calls(self, tool_calls: list[ChatCompletionDeltaToolCall]) -> None:
+        """Handle tool call updates"""
+        if not tool_calls:
+            return
+        for call in tool_calls:
+            if call.id:
+                new_tool_call = ToolCall(
+                    id=call.id,
+                    type=call.type,
+                    function=ToolCallFunction(
+                        name=call.function.name or "",
+                        arguments=call.function.arguments,
+                    ),
+                    index=call.index,
+                )
+                if self.current_message.tool_calls is None:
+                    self.current_message.tool_calls = []
+                self.current_message.tool_calls.append(new_tool_call)
+            else:
+                existing_call = self.current_message.tool_calls[call.index]
+                if call.function.arguments:
+                    existing_call.function.arguments += call.function.arguments
+
+    def handle_usage_info(self, chunk: litellm.ModelResponseStream) -> litellm.Usage | None:
+        """Handle usage info, return whether this chunk should be skipped"""
+        usage = getattr(chunk, "usage", None)
+        if usage:
+            logger.debug("Model usage: %s", usage)
+        return usage
+
+    def finalize_message(self) -> AssistantMessage:
+        """Finalize message processing"""
+        logger.debug("Message finalized: %s", self.current_message)
+        return self.current_message

lite_agent/runner.py

@@ -0,0 +1,51 @@
+from collections.abc import AsyncGenerator
+from typing import Literal
+
+from open_agents.agent import Agent
+from open_agents.chunk_handler import AgentChunk
+from open_agents.types import AgentToolCallMessage, RunnerMessages
+
+
+class Runner:
+    def __init__(self, agent: Agent) -> None:
+        self.agent = agent
+        self.messages: RunnerMessages = []
+
+    def run_stream(
+        self,
+        user_input: RunnerMessages | str,
+        max_steps: int = 20,
+        includes: list[Literal["usage", "final_message", "tool_call", "tool_call_result"]] | None = None,
+    ) -> AsyncGenerator[AgentChunk, None]:
+        """Run the agent and return a RunResponse object that can be asynchronously iterated for each chunk."""
+        if includes is None:
+            includes = ["final_message", "usage", "tool_call", "tool_call_result", "tool_call_delta", "content_delta"]
+        if isinstance(user_input, str):
+            self.messages.append({"role": "user", "content": user_input})
+        else:
+            self.messages = user_input
+
+        return self._run_aiter(max_steps, includes)
+
+    async def _run_aiter(self, max_steps: int, includes: list[Literal["usage", "final_message", "tool_call", "tool_call_result"]]) -> AsyncGenerator[AgentChunk, None]:
+        """Run the agent and return a RunResponse object that can be asynchronously iterated for each chunk."""
+        steps = 0
+        finish_reason = None
+        while finish_reason != "stop" and steps < max_steps:
+            resp = await self.agent.stream_async(self.messages)
+            async for chunk in resp:
+                if chunk["type"] == "final_message":
+                    message = chunk["message"]
+                    self.messages.append(message.model_dump())
+                    finish_reason = chunk["finish_reason"]
+                elif chunk["type"] == "tool_call_result":
+                    self.messages.append(
+                        AgentToolCallMessage(
+                            role="tool",
+                            tool_call_id=chunk["tool_call_id"],
+                            content=chunk["content"],
+                        ),
+                    )
+                if chunk["type"] in includes:
+                    yield chunk
+            steps += 1

lite_agent/types.py

@@ -0,0 +1,152 @@
+from typing import Literal, TypedDict
+
+import litellm
+from pydantic import BaseModel
+
+
+class ToolCallFunction(BaseModel):
+    name: str
+    arguments: str | None = None
+
+
+class ToolCall(BaseModel):
+    type: Literal["function"]
+    function: ToolCallFunction
+    id: str
+
+
+class AssistantMessage(BaseModel):
+    id: str
+    role: Literal["assistant"] = "assistant"
+    content: str = ""
+    tool_calls: list[ToolCall] | None = None
+
+
+class Message(TypedDict):
+    role: str
+    content: str
+
+
+class UserMessageContentItemText(TypedDict):
+    type: Literal["text"]
+    text: str
+
+
+class UserMessageContentItemImageURLImageURL(TypedDict):
+    url: str
+
+
+class UserMessageContentItemImageURL(TypedDict):
+    type: Literal["image_url"]
+    image_url: UserMessageContentItemImageURLImageURL
+
+
+class AgentUserMessage(TypedDict):
+    role: Literal["user"] = "user"
+    content: str | list[UserMessageContentItemText | UserMessageContentItemImageURL]
+
+
+class AssistantMessageToolCallFunction(TypedDict):
+    name: str
+    arguments: str
+
+
+class AssistantMessageToolCall(TypedDict):
+    id: str
+    type: Literal["function"]
+    function: AssistantMessageToolCallFunction
+    tool_call_id: str
+
+
+class AgentAssistantMessage(TypedDict):
+    role: Literal["assistant"] = "assistant"
+    content: str
+    tool_calls: list[AssistantMessageToolCall] | None
+
+
+class AgentSystemMessage(TypedDict):
+    role: Literal["system"] = "system"
+    content: str
+
+
+class AgentToolCallMessage(TypedDict):
+    role: Literal["tool"] = "tool"
+    tool_call_id: str
+    content: str
+
+
+RunnerMessage = AgentUserMessage | AgentAssistantMessage | AgentToolCallMessage
+AgentMessage = RunnerMessage | AgentSystemMessage
+RunnerMessages = list[RunnerMessage]
+
+
+class LiteLLMRawChunk(TypedDict):
+    """
+    Define the type of chunk
+    """
+
+    type: Literal["litellm_raw"]
+    raw: litellm.ModelResponseStream
+
+
+class UsageChunk(TypedDict):
+    """
+    Define the type of usage info chunk
+    """
+
+    type: Literal["usage"]
+    usage: litellm.Usage
+
+
+class FinalMessageChunk(TypedDict):
+    """
+    Define the type of final message chunk
+    """
+
+    type: Literal["final_message"]
+    message: AssistantMessage
+    finish_reason: Literal["stop", "tool_calls"]
+
+
+class ToolCallChunk(TypedDict):
+    """
+    Define the type of tool call chunk
+    """
+
+    type: Literal["tool_call"]
+    name: str
+    arguments: str
+
+
+class ToolCallResultChunk(TypedDict):
+    """
+    Define the type of tool call result chunk
+    """
+
+    type: Literal["tool_call_result"]
+    tool_call_id: str
+    name: str
+    content: str
+
+
+class ContentDeltaChunk(TypedDict):
+    """
+    Define the type of message chunk
+    """
+
+    type: Literal["content_delta"]
+    delta: str
+
+
+class ToolCallDeltaChunk(TypedDict):
+    """
+    Define the type of tool call delta chunk
+    """
+
+    type: Literal["tool_call_delta"]
+    tool_call_id: str
+    name: str
+    arguments_delta: str
+
+
+AgentChunk = LiteLLMRawChunk | UsageChunk | FinalMessageChunk | ToolCallChunk | ToolCallResultChunk | ContentDeltaChunk | ToolCallDeltaChunk

lite_agent-0.1.0.dist-info/METADATA

@@ -0,0 +1,22 @@
+Metadata-Version: 2.4
+Name: lite-agent
+Version: 0.1.0
+Summary: A lightweight, extensible framework for building AI agent.
+Author-email: Jianqi Pan <jannchie@gmail.com>
+License: MIT
+Keywords: AI,agent framework,assistant,chatbot,function call,openai,pydantic,rich,tooling
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Communications :: Chat
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.10
+Requires-Dist: funcall>=0.6.0
+Requires-Dist: prompt-toolkit>=3.0.51
+Requires-Dist: rich>=14.0.0

lite_agent-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+lite_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lite_agent/__main__.py,sha256=D4Yg5a6hDYMq6aCEGg3MjILF4pVWD6mBKjri-qs7Rzk,3766
+lite_agent/agent.py,sha256=LmC5aqMVWsq5AkE18Od0FEVUdiZN1GTs5OlO6mrvxUM,1187
+lite_agent/chunk_handler.py,sha256=MDk9rW_axZ412VOPnITbQGVr-mCGAe2jy3khUXkYbvM,5135
+lite_agent/loggers.py,sha256=0-yb_1Ec8TodfTI_nZWeCJSpSZZM9R53UUOE5--L2iw,57
+lite_agent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lite_agent/runner.py,sha256=TaBm6pCFwdJtygNnmVZDlFbkK3-Z1s52YD9XviwrA6o,2257
+lite_agent/types.py,sha256=NUpoSDaaiWWYezWluDdWBb-DsYFvZFoNd4tLWqBzCZ4,3058
+lite_agent/processors/__init__.py,sha256=AIxxS9uLcNws1Erf-b8z8zZKaKA5NQ7ooZ7HaCIJTq4,115
+lite_agent/processors/stream_chunk_processor.py,sha256=WlVcfc4kdiCz0Fc1Ybe5GapN2stK0C_0DIymbiGpZfI,3420
+lite_agent-0.1.0.dist-info/METADATA,sha256=oFRlVtYwsr1P9RPOP2JNu309-Jpqc95f-TYrExmrFBc,972
+lite_agent-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+lite_agent-0.1.0.dist-info/RECORD,,

lite_agent-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any