Add sources

From private repo commit 55315374

Author: rivol

.gitignore

@@ -0,0 +1,6 @@
+/.idea/
+/.venv/
+/dist/
+__pycache__/
+
+/.env

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Rivo Laks
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,93 @@
+# AAF
+
+AAF is a powerful and flexible framework for building and managing conversational AI models.
+It provides a unified interface for various language model providers and
+implements advanced virtual models for complex conversational scenarios.
+
+
+## Features
+
+- Support for multiple LLM providers (OpenAI, Anthropic, Ollama)
+- Advanced conversation management with Threads and Sessions
+- Virtual models for complex scenarios (TwoPhase, Multiphase, Router)
+- Tool integration for function calling capabilities
+- Cost and token usage tracking
+
+
+## Installation
+
+(Add installation instructions here)
+
+
+## Quick Start
+
+```python
+from aaf.threads import Session
+
+thread = Session().create_thread("gpt-4o", system="You are a helpful assistant.")
+thread.add_message("user", "What is the capital of France?")
+
+async with thread.run() as stream:
+    async for chunk in stream.text_chunks():
+        print(chunk.content, end="", flush=True)
+    print()
+
+print(thread.cost_and_usage().pretty())
+```
+
+
+## LLM Providers
+
+AAF supports the following LLM providers:
+- OpenAI
+- Anthropic
+- Ollama
+
+To use a specific provider, specify the model name when creating a thread:
+
+```python
+thread = session.create_thread("gpt-4o")  # OpenAI
+thread = session.create_thread("claude-3-5-sonnet-20240620")  # Anthropic
+thread = session.create_thread("llama3:instruct")  # Ollama
+```
+
+
+## Virtual Models
+
+AAF implements several virtual models for advanced use cases:
+
+- TwoPhase: Generates a prompt and then uses it to create a response
+- Multiphase: Multi-step process for complex questions, including drafting, feedback, and refinement
+- Router: Selects the appropriate model based on the user's request
+
+Using a virtual model is same as with standard models:
+
+```python
+from aaf.virtual_models.two_phase import TwoPhaseModel
+from aaf.threads import Session
+
+thread = Session().create_thread(model="two-phase", runner=TwoPhaseModel())
+thread.add_message("user", "What is the capital of France?")
+
+async with thread.run() as stream:
+    async for chunk in stream.text_chunks():
+        print(chunk.content, end="", flush=True)
+    print()
+
+print(thread.cost_and_usage().pretty())
+```
+
+
+## Project Structure
+
+- `aaf/`: Main package directory
+  - `llms/`: LLM provider implementations
+  - `virtual_models/`: Virtual model implementations
+  - `threads.py`: Thread and Session management
+  - `logging.py`: Custom logging implementation
+  - `utils.py`: Utility functions
+
+
+## License
+
+(Add license information)

pyproject.toml

@@ -1,10 +1,22 @@
 [project]
 name = "aaf"
-version = "0.1.0"
+version = "0.3.0"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12"
-dependencies = []
+dependencies = [
+    "anthropic~=0.29.0",
+    "docstring-parser~=0.16",
+    "duckduckgo-search~=6.2.6",
+    "fastapi~=0.111.0",
+    "openai~=1.30.1",
+    "prompt-toolkit~=3.0.47",
+    "pydantic==2.7.1",
+    "rich~=13.7.1",
+    "sse-starlette~=2.1.2",
+    "typer~=0.12.3",
+    "typing-inspect~=0.9.0",
+]
 
 [build-system]
 requires = ["hatchling"]

src/aaf/__init__.py

@@ -1,2 +0,0 @@
-def hello() -> str:
-    return "Hello from aaf!"

src/aaf/chat.py

@@ -0,0 +1,59 @@
+from typing import Optional
+
+from prompt_toolkit import PromptSession
+from prompt_toolkit.patch_stdout import patch_stdout
+
+from aaf.event_handlers import use_event_handler
+from aaf.threads import Thread
+from aaf.ui import SimpleStreamingUserInterface
+
+
+class ChatSession:
+    """
+    Manages an interactive chat session with an AI model.
+
+    This class handles the user input loop, processes questions through the AI model,
+    and manages the overall flow of the conversation, including handling user exit commands.
+
+    Attributes:
+        thread (Thread): The Thread object managing the conversation with the AI model.
+        ui (SimpleStreamingUserInterface): The user interface for displaying output and handling input.
+    """
+
+    def __init__(self, thread: Thread, ui: Optional[SimpleStreamingUserInterface] = None):
+        self.thread = thread
+        self.ui = ui or SimpleStreamingUserInterface()
+        self.prompt_session = PromptSession()
+
+    async def run_loop(self, initial_question: Optional[str] = None, *, interactive: bool = True):
+        if initial_question:
+            await self._process_question(initial_question)
+
+        try:
+            while interactive:
+                question = await self._get_user_input()
+                if question is None:
+                    break
+                await self._process_question(question)
+
+        finally:
+            self.ui.info(f"Final cost: {self.thread.cost_and_usage().pretty_root()}")
+
+    async def _get_user_input(self) -> Optional[str]:
+        with patch_stdout():
+            try:
+                question = await self.prompt_session.prompt_async("You: ")
+                if question.lower() in ("exit", "quit"):
+                    return None
+                return question
+            except EOFError:
+                # Handle Ctrl+D
+                print("\nGoodbye!")
+                return None
+
+    async def _process_question(self, question: str):
+        self.thread.add_message("user", question)
+        async with self.thread.run_loop() as stream:
+            await use_event_handler(stream, self.ui)
+
+        self.ui.info(f"Current cost: {self.thread.cost_and_usage().pretty_root()}")

src/aaf/event_handlers.py

@@ -0,0 +1,73 @@
+from .llms.types import (
+    ResponseChunkToolCallFailed,
+    ResponseChunkToolCallFinished,
+    ResponseChunkToolCallStarted,
+    ResponseControlChunkStreamBegin,
+    ResponseControlChunkStreamEnd,
+    ResponseStream,
+    ResponseTextChunk,
+    ToolCall,
+)
+
+
+class EventHandlerBase:
+    """Base class for event handlers that process messages and events from the assistant and tools.
+
+    This can be used to implement User Interface classes, custom logic, etc.
+    """
+
+    def info(self, message: str):
+        """Print an informational message."""
+        pass
+
+    def debug(self, message: str):
+        """Print a debug message."""
+        pass
+
+    def assistant_message_stream_start(self, stream: ResponseStream):
+        """Called when the assistant starts sending a message stream."""
+        pass
+
+    def assistant_message_stream_chunk(self, stream: ResponseStream, chunk: ResponseTextChunk):
+        """Called for each chunk of the assistant's response."""
+        pass
+
+    def assistant_message_stream_end(self, stream: ResponseStream):
+        """Called when the assistant finishes sending a message stream."""
+        pass
+
+    def tool_call_start(self, tool_call: ToolCall):
+        """Called when a tool call is about to be executed."""
+        pass
+
+    def tool_call_result(self, tool_call: ToolCall, result: str):
+        """Called when a tool call has been executed successfully."""
+        pass
+
+    def tool_call_error(self, tool_call: ToolCall, error: str):
+        """Called when a tool call has failed."""
+        pass
+
+    def loop_end(self, cost: float):
+        pass
+
+
+async def use_event_handler(stream: ResponseStream, handler: EventHandlerBase):
+    """Uses the event handler for the stream's messages.
+
+    This is basically an adapter that directs the stream's messages to the event handler.
+    """
+
+    async for chunk in stream.all_chunks():
+        if isinstance(chunk, ResponseControlChunkStreamBegin):
+            handler.assistant_message_stream_start(stream)
+        elif isinstance(chunk, ResponseControlChunkStreamEnd):
+            handler.assistant_message_stream_end(stream)
+        elif isinstance(chunk, ResponseTextChunk):
+            handler.assistant_message_stream_chunk(stream, chunk)
+        elif isinstance(chunk, ResponseChunkToolCallStarted):
+            handler.tool_call_start(chunk.tool_call)
+        elif isinstance(chunk, ResponseChunkToolCallFinished):
+            handler.tool_call_result(chunk.tool_call, chunk.result)
+        elif isinstance(chunk, ResponseChunkToolCallFailed):
+            handler.tool_call_error(chunk.tool_call, chunk.error)

src/aaf/forwarding.py

@@ -0,0 +1,73 @@
+import asyncio
+from asyncio import Queue
+from typing import AsyncIterator
+
+from .llms.base import ResponseAdapterBase
+from .llms.types import (
+    ResponseChunk,
+    ResponseDebugChunk,
+    ResponseStopReasonChunk,
+    ResponseTextChunk,
+    ResponseVerboseChunk,
+    StopReason,
+)
+
+
+class ResponseQueue(Queue):
+    """Queue for sending response chunks.
+
+    This is similar to ResponseStream, but focuses on the writing part (whereas ResponseStream focuses on reading).
+    It's a subclass of asyncio.Queue, with some convenience methods for adding different types of chunks.
+    """
+
+    async def add(self, content: str) -> None:
+        await self.put(ResponseTextChunk(content=content))
+
+    async def add_debug(self, content: str) -> None:
+        await self.put(ResponseDebugChunk(content=content))
+
+    async def add_verbose(self, content: str) -> None:
+        await self.put(ResponseVerboseChunk(content=content))
+
+    async def mark_finished(self) -> None:
+        await self.put(None)
+
+
+class VirtualModelAdapter(ResponseAdapterBase):
+    """Response stream adapter for virtual models.
+
+    It reads chunks from the provided queue and forwards them to the stream.
+    The queue must end with a None value to indicate the end of the stream.
+    """
+
+    def __init__(self, process_task: asyncio.Task[None], queue: ResponseQueue):
+        super().__init__()
+
+        self.process_task = process_task
+        self.queue = queue
+
+        self._stop_reason = None
+
+    async def __aiter__(self) -> AsyncIterator[ResponseChunk]:
+        """Forward chunks from the queue to the stream.
+
+        Takes care to avoid yielding the stop reason twice.
+        """
+
+        try:
+            while True:
+                chunk = await self.queue.get()
+                if chunk is None:  # End of stream marker
+                    if self._stop_reason is None:
+                        yield ResponseStopReasonChunk(reason=StopReason.END_TURN)
+
+                    break
+
+                if isinstance(chunk, ResponseStopReasonChunk):
+                    self._stop_reason = chunk.reason
+
+                yield chunk
+
+        finally:
+            # Ensure the process task is completed
+            await self.process_task

src/aaf/llms/__init__.py

@@ -0,0 +1,10 @@
+from .anthropic import AnthropicRunner
+from .base import ModelRunner
+from .ollama import OllamaRunner
+from .openai import OpenAIRunner
+
+PROVIDERS = [OpenAIRunner, AnthropicRunner, OllamaRunner]
+
+
+def get_llm_provider_and_model(model_name: str) -> tuple[type["ModelRunner"], str]:
+    return ModelRunner.get_provider_and_model(PROVIDERS, model_name)