psi-agent 0.0.1a2__py3-none-any.whl

psi_agent/__init__.py

@@ -0,0 +1 @@
+"""Psi Agent Platform - A portable, composable AI Agent platform."""

psi_agent/ai/__init__.py

@@ -0,0 +1,5 @@
+"""Psi AI - LLM Caller interfaces."""
+
+from . import openai
+
+__all__ = ["openai"]

psi_agent/ai/openai/__init__.py

@@ -0,0 +1,197 @@
+"""Psi AI OpenAI - LLM Caller that exposes OpenAI-compatible API via Unix socket."""
+
+import asyncio
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import tyro
+from loguru import logger
+from openai import AsyncOpenAI
+
+from psi_agent.common import LLMRequest, LLMResponse
+
+
+class AICaller:
+    """LLM Caller that forwards requests to OpenAI-compatible APIs."""
+
+    def __init__(
+        self,
+        session_socket: str,
+        api_key: str,
+        base_url: str,
+        model: str,
+    ) -> None:
+        self._session_socket = session_socket
+        self._client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+        self._model = model
+
+        logger.info(f"AI Caller initialized | model={model} | base_url={base_url}")
+        logger.debug(f"AI Caller config | socket={session_socket}")
+
+    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        """Handle a single client connection."""
+        data = await reader.readline()
+        if not data:
+            logger.debug("Empty request received")
+            return
+
+        request = LLMRequest.model_validate(json.loads(data.decode()))
+        logger.info(
+            f"Request received | id={request.id} | stream={request.stream} | message_count={len(request.messages)}"
+        )
+        logger.debug(f"Request details | has_tools={request.tools is not None} | tool_choice={request.tool_choice}")
+
+        if request.stream:
+            await self._handle_stream(request, writer)
+        else:
+            await self._handle_non_stream(request, writer)
+
+        writer.close()
+        await writer.wait_closed()
+
+    async def _handle_stream(self, request: LLMRequest, writer: asyncio.StreamWriter) -> None:
+        """Handle streaming request."""
+        logger.debug(f"Starting stream | request_id={request.id}")
+
+        kwargs: dict[str, Any] = {
+            "model": self._model,
+            "messages": request.messages,
+            "stream": True,
+        }
+        if request.tools:
+            kwargs["tools"] = request.tools
+            kwargs["tool_choice"] = request.tool_choice
+
+        stream = await self._client.chat.completions.create(**kwargs)
+        chunk_count = 0
+
+        async for chunk in stream:
+            chunk_count += 1
+            chunk_data = chunk.model_dump()
+            response = LLMResponse(id=request.id, choices=chunk_data.get("choices", []))
+            writer.write((response.model_dump_json() + "\n").encode())
+            await writer.drain()
+
+        # Send done marker
+        done_response = LLMResponse(id=request.id, choices=[], done=True)
+        writer.write((done_response.model_dump_json() + "\n").encode())
+        await writer.drain()
+
+        logger.info(f"Stream complete | request_id={request.id} | chunks={chunk_count}")
+
+    async def _handle_non_stream(self, request: LLMRequest, writer: asyncio.StreamWriter) -> None:
+        """Handle non-streaming request."""
+        logger.debug(f"Starting non-stream request | request_id={request.id}")
+
+        kwargs: dict[str, Any] = {
+            "model": self._model,
+            "messages": request.messages,
+        }
+        if request.tools:
+            kwargs["tools"] = request.tools
+            kwargs["tool_choice"] = request.tool_choice
+
+        response = await self._client.chat.completions.create(**kwargs)
+        response_data = response.model_dump()
+        llm_response = LLMResponse(id=request.id, choices=response_data.get("choices", []))
+        writer.write((llm_response.model_dump_json() + "\n").encode())
+        await writer.drain()
+
+        content = response_data.get("choices", [{}])[0].get("message", {}).get("content", "")
+        logger.info(f"Non-stream complete | request_id={request.id} | response_length={len(content)}")
+
+    async def run(self) -> None:
+        """Start the Unix socket server."""
+        socket_path = Path(self._session_socket)
+        if socket_path.exists():
+            socket_path.unlink()
+            logger.debug(f"Removed existing socket | path={self._session_socket}")
+
+        server = await asyncio.start_unix_server(
+            self.handle_client,
+            path=self._session_socket,
+        )
+
+        logger.info(f"AI server started | socket={self._session_socket} | model={self._model}")
+
+        async with server:
+            await server.serve_forever()
+
+
+async def run_ai(
+    session_socket: str,
+    model: str,
+    api_key: str,
+    base_url: str = "https://api.openai.com/v1",
+    log_level: str = "INFO",
+) -> None:
+    """Python function interface to run the AI caller.
+
+    Args:
+        session_socket: Unix socket path to listen on
+        model: Model name to use
+        api_key: API key for the LLM provider
+        base_url: API base URL
+        log_level: Log level (DEBUG, INFO, WARNING, ERROR)
+    """
+    _setup_logger(log_level)
+    caller = AICaller(
+        session_socket=session_socket,
+        api_key=api_key,
+        base_url=base_url,
+        model=model,
+    )
+    await caller.run()
+
+
+def _setup_logger(log_level: str) -> None:
+    """Configure logger."""
+    logger.remove()
+    logger.add(
+        lambda msg: print(msg, end=""),
+        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>ai-openai</cyan> | {message}",
+        level=log_level,
+    )
+
+
+@dataclass
+class CliArgs:
+    """AI Caller CLI arguments."""
+
+    session_socket: str
+    """Unix socket path to listen on"""
+
+    model: str
+    """Model name to use"""
+
+    api_key: str | None = None
+    """API key (or set API_KEY env var)"""
+
+    base_url: str = "https://api.openai.com/v1"
+    """API base URL"""
+
+    log_level: str = "INFO"
+    """Log level (DEBUG, INFO, WARNING, ERROR)"""
+
+
+def main() -> None:
+    args = tyro.cli(CliArgs)
+
+    api_key = args.api_key or os.environ.get("API_KEY") or ""
+
+    asyncio.run(
+        run_ai(
+            session_socket=args.session_socket,
+            model=args.model,
+            api_key=api_key,
+            base_url=args.base_url,
+            log_level=args.log_level,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

psi_agent/channel/__init__.py

@@ -0,0 +1,5 @@
+"""Psi Channel - User interaction interfaces."""
+
+from . import tui
+
+__all__ = ["tui"]

psi_agent/channel/tui/__init__.py

@@ -0,0 +1,117 @@
+"""Psi Channel - TUI interface for user interaction."""
+
+import asyncio
+import json
+import sys
+from dataclasses import dataclass
+
+import tyro
+from loguru import logger
+from prompt_toolkit import PromptSession
+from prompt_toolkit.styles import Style
+
+from psi_agent.common import AssistantMessage, UserMessage
+
+
+class Channel:
+    """Terminal UI channel for user interaction."""
+
+    def __init__(self, session_socket: str) -> None:
+        self._session_socket = session_socket
+        logger.debug(f"Channel initialized | socket={session_socket}")
+
+    async def run(self) -> None:
+        """Run the TUI interface."""
+        logger.info(f"Connecting to session at {self._session_socket}")
+        print(f"Connecting to session at {self._session_socket}", file=sys.stderr)
+
+        reader, writer = await asyncio.open_unix_connection(self._session_socket)
+        logger.info("Connected to session")
+
+        style = Style.from_dict(
+            {
+                "user": "#ansigreen",
+                "assistant": "#ansicyan",
+            }
+        )
+
+        prompt_session: PromptSession[str] = PromptSession(style=style)
+
+        print("Connected. Ctrl+C to exit.")
+
+        while True:
+            try:
+                user_input = await prompt_session.prompt_async("You: ", style=style)
+                logger.debug(f"User input received | length={len(user_input)}")
+
+                if not user_input:
+                    continue
+
+                message = UserMessage(content=user_input)
+                writer.write((message.model_dump_json() + "\n").encode())
+                await writer.drain()
+                logger.debug(f"Message sent to session | content={user_input[:50]}")
+
+                data = await reader.readline()
+                if not data:
+                    logger.warning("Session disconnected")
+                    print("Session disconnected.")
+                    break
+
+                response = AssistantMessage.model_validate(json.loads(data.decode()))
+                logger.debug(f"Response received | length={len(response.content)}")
+                print(f"Assistant: {response.content}")
+
+            except KeyboardInterrupt:
+                logger.info("User interrupted with Ctrl+C")
+                print("\nExiting...")
+                break
+            except EOFError:
+                logger.info("EOF received")
+                break
+
+        writer.close()
+        await writer.wait_closed()
+        logger.info("Channel closed")
+
+
+async def run_channel(session_socket: str, log_level: str = "WARNING") -> None:
+    """Python function interface to run the TUI channel.
+
+    Args:
+        session_socket: Unix socket path for session connection
+        log_level: Log level (DEBUG, INFO, WARNING, ERROR). Default WARNING to not show logs in TUI.
+    """
+    _setup_logger(log_level)
+    channel = Channel(session_socket=session_socket)
+    await channel.run()
+
+
+def _setup_logger(log_level: str) -> None:
+    """Configure logger."""
+    logger.remove()
+    logger.add(
+        lambda msg: print(msg, end=""),
+        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>channel-tui</cyan> | {message}",
+        level=log_level,
+    )
+
+
+@dataclass
+class CliArgs:
+    """Channel TUI CLI arguments."""
+
+    session_socket: str
+    """Unix socket path for session connection"""
+
+    log_level: str = "WARNING"
+    """Log level (DEBUG, INFO, WARNING, ERROR). Default WARNING to not show logs in TUI"""
+
+
+def main() -> None:
+    args = tyro.cli(CliArgs)
+    asyncio.run(run_channel(session_socket=args.session_socket, log_level=args.log_level))
+
+
+if __name__ == "__main__":
+    main()

psi_agent/common/__init__.py

@@ -0,0 +1,5 @@
+"""Psi common utilities."""
+
+from .protocol import AssistantMessage, LLMRequest, LLMResponse, ToolResult, UserMessage
+
+__all__ = ["LLMRequest", "LLMResponse", "ToolResult", "UserMessage", "AssistantMessage"]

psi_agent/common/protocol.py

@@ -0,0 +1,49 @@
+"""Shared protocol definitions for Psi Agent Platform."""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class LLMRequest(BaseModel):
+    """Request to LLM Caller (OpenAI compatible)."""
+
+    id: str
+    messages: list[dict[str, Any]]
+    tools: list[dict[str, Any]] | None = None
+    tool_choice: str = "auto"
+    stream: bool = True
+
+
+class LLMResponse(BaseModel):
+    """Response from LLM Caller (OpenAI compatible)."""
+
+    id: str
+    choices: list[dict[str, Any]]
+    done: bool = False
+    error: str | None = None
+
+
+class ToolResult(BaseModel):
+    """Result from tool execution."""
+
+    success: bool
+    content: str | None = None
+    error: str | None = None
+    stdout: str | None = None
+    stderr: str | None = None
+    returncode: int | None = None
+
+
+class UserMessage(BaseModel):
+    """Message from user to session."""
+
+    role: str = "user"
+    content: str
+
+
+class AssistantMessage(BaseModel):
+    """Message from assistant to channel."""
+
+    role: str = "assistant"
+    content: str