From 056a1a0f47a7e78a33b2494889c4382e4de939ea Mon Sep 17 00:00:00 2001 From: prasangeet Date: Tue, 25 Aug 2026 21:19:46 +0530 Subject: [PATCH 1/5] feat: major changes, updated mcp architecture and added multi-llm integration --- README.md | 534 +++++++----- client/src/app.rs | 283 +++++- client/src/ipc/events.rs | 85 +- client/src/ipc/messages.rs | 274 +++++- runtime/mcp.json | 12 +- runtime/pyproject.toml | 1 + runtime/src/orion/agent/nodes/remember.py | 140 ++- runtime/src/orion/agent/prompts.py | 92 +- runtime/src/orion/events/base.py | 4 +- runtime/src/orion/integrations/_mcp/config.py | 169 +++- runtime/src/orion/integrations/_mcp/events.py | 127 +++ .../src/orion/integrations/_mcp/langchain.py | 123 +++ .../integrations/_mcp/langchain_tools.py | 113 --- .../src/orion/integrations/_mcp/manager.py | 466 +++++++++- runtime/src/orion/integrations/_mcp/server.py | 168 +++- runtime/src/orion/llm/__init__.py | 0 runtime/src/orion/llm/config.py | 57 ++ runtime/src/orion/llm/factory.py | 45 + runtime/src/orion/llm/interfaces/provider.py | 16 + runtime/src/orion/llm/providers/gemini.py | 35 + runtime/src/orion/llm/providers/groq.py | 37 + runtime/src/orion/llm/utils.py | 33 + runtime/src/orion/memory/config.py | 85 +- runtime/src/orion/memory/interfaces/graph.py | 79 +- runtime/src/orion/memory/models.py | 16 +- .../src/orion/memory/providers/graph/neo4j.py | 254 ++++-- runtime/src/orion/memory/session.py | 66 +- runtime/src/orion/orchestrator/config.py | 6 +- .../src/orion/orchestrator/orchestrator.py | 5 + runtime/src/orion/runtime/run.py | 29 +- runtime/src/orion/runtime/runtime.py | 4 + runtime/src/orion/services/agent.py | 70 +- runtime/src/orion/services/base.py | 62 +- runtime/src/orion/services/ipc_publisher.py | 3 +- runtime/src/orion/services/setup.py | 11 +- .../orion/services/transcript_generation.py | 298 +++---- runtime/src/orion/transport/bridge.py | 2 + runtime/src/orion/transport/messages.py | 198 ++++- .../tests/integrations/test_mcp_discovery.py | 183 ++++ .../tests/integrations/test_mcp_langchain.py | 378 ++++++++ .../tests/integrations/test_mcp_manager.py | 806 ++++++++++++++++++ runtime/tests/integrations/test_mcp_server.py | 439 ++++++++++ runtime/tests/llm/test_gemini.py | 62 ++ runtime/tests/llm/test_groq.py | 62 ++ runtime/tests/memory/test_neo4j.py | 777 +++++++++++++++++ runtime/uv.lock | 93 +- 46 files changed, 5900 insertions(+), 902 deletions(-) create mode 100644 runtime/src/orion/integrations/_mcp/events.py create mode 100644 runtime/src/orion/integrations/_mcp/langchain.py delete mode 100644 runtime/src/orion/integrations/_mcp/langchain_tools.py create mode 100644 runtime/src/orion/llm/__init__.py create mode 100644 runtime/src/orion/llm/config.py create mode 100644 runtime/src/orion/llm/factory.py create mode 100644 runtime/src/orion/llm/interfaces/provider.py create mode 100644 runtime/src/orion/llm/providers/gemini.py create mode 100644 runtime/src/orion/llm/providers/groq.py create mode 100644 runtime/src/orion/llm/utils.py create mode 100644 runtime/tests/integrations/test_mcp_discovery.py create mode 100644 runtime/tests/integrations/test_mcp_langchain.py create mode 100644 runtime/tests/integrations/test_mcp_manager.py create mode 100644 runtime/tests/integrations/test_mcp_server.py create mode 100644 runtime/tests/llm/test_gemini.py create mode 100644 runtime/tests/llm/test_groq.py create mode 100644 runtime/tests/memory/test_neo4j.py diff --git a/README.md b/README.md index 23d02a7..d380abd 100644 --- a/README.md +++ b/README.md @@ -5,316 +5,448 @@

- Python 3.14+ + Python 3.11+ + Rust edition 2024 Event Driven - Textual UI - Voice Native + LangGraph + Qdrant + Neo4j

ORION is a unified intelligent system designed to combine reasoning, perception, memory, planning, and execution into a single event-driven platform. -It is not just a chatbot. The codebase already implements a voice-first pipeline that records audio, transcribes speech, generates responses with an LLM, synthesizes speech, plays it back, and persists the full event trail in SQLite. The longer-term goal is to grow this into a modular intelligence platform that can understand goals, manage context, coordinate tools, and eventually control a computer and assist with coding tasks in the style of modern agentic systems. +It is not just a chatbot. The codebase implements a memory-aware agent runtime that runs as a background daemon, a terminal client that talks to it over a Unix socket, a LangGraph agent that retrieves long-term memory and calls both local and MCP tools, and an event store that records every step. The longer-term goal is to grow this into a modular intelligence platform that can understand goals, manage context, coordinate tools, and eventually control a computer and assist with coding tasks in the style of modern agentic systems. -## What ORION Does Today +## What ORION Is Today -The current implementation is centered around a strict event pipeline: +ORION runs as **two processes** that communicate over a Unix domain socket: -1. Start a pipeline. -2. Record microphone input. -3. Transcribe the recording with Groq Whisper. -4. Generate a response with Groq chat completions. -5. Synthesize speech with Kokoro. -6. Play the generated audio back through the speakers. -7. Persist every event into SQLite. -8. Render the live stream in a Textual terminal UI. +- **`runtime/`** — a Python daemon that owns the event bus, the agent, memory, MCP servers, and the event store. It never touches audio hardware. +- **`client/`** — a Rust (Ratatui) terminal client that owns the terminal, the microphone, and speech output. + +Everything crossing the boundary is a newline-delimited JSON envelope with a `type` and a `payload`. + +### How a Request Flows + +1. The client connects to `/tmp/orion.sock` and renders the TUI. +2. You type a prompt (`i` for insert mode, `Enter` to send) or press `v` to record your voice. +3. For voice, the client captures the microphone to a 16-bit PCM WAV and sends the runtime the file path; the runtime transcribes it with Groq Whisper. +4. Either path publishes a `ChatPipelineStartEvent` onto the event bus. +5. The `AgentService` builds a LangGraph graph and invokes it: retrieve memory → reason → call tools → loop → remember. +6. The response is published as `assistant_start` / `assistant_chunk` / `assistant_end` messages over IPC. +7. The client streams the reply into the conversation panel and speaks it through the system speech engine. +8. Every event is appended to SQLite on the way through the bus, and mirrored into the client's live event stream panel. ### Current Capabilities -- Event-driven orchestration with a shared bus -- Persistent event storage in SQLite -- Voice recording via `sounddevice` -- Speech-to-text via Groq -- Response generation via Groq chat completions -- Text-to-speech via Kokoro -- Audio playback via `sounddevice` and `soundfile` -- Live terminal interface via Textual -- CLI entrypoint via Typer +- Event-driven orchestration with a shared bus and a strict lifecycle (`OrionRuntime` → components → services) +- Persistent event storage in SQLite, written before any handler runs +- Pluggable LLM provider layer — Google Gemini (default) or Groq, selected by environment variable +- LangGraph agent with a `retrieve → agent → tools → remember` loop +- Long-term memory: Qdrant vector recall, a Neo4j knowledge graph of extracted facts, a SQLite rolling summary, and local sentence-transformer embeddings +- Automatic fact extraction after each turn, written to the knowledge graph +- MCP client manager — starts configured servers, discovers their tools, routes calls, and publishes lifecycle/execution events +- Local automation tools (open browser/URL/terminal/file manager, launch applications, run shell commands) +- IPC transport — Unix socket server, NDJSON protocol, typed Pydantic envelopes, per-client sessions +- Speech-to-text via Groq Whisper (`whisper-large-v3-turbo`) +- Text-to-speech synthesis via Kokoro +- Rust terminal client with vim-style keybindings, mouse support, a live event stream, Copilot-style activity logs, and `tachyonfx` animations +- Microphone capture in the client via `cpal` + `hound`; spoken replies via `spd-say` ### Current CLI Surface -- `orion voice` - launches the live voice pipeline and TUI -- `orion chat` - placeholder for future chat mode -- `orion doctor` - basic environment/status check +The runtime has a single entrypoint — running it starts the daemon and blocks until shutdown: + +```bash +uv run orion # or: uv run python -m orion +``` + +There are no subcommands. The client is a separate binary (`cargo run` in `client/`). ## Architecture -ORION is intentionally structured as a pipeline of small services coordinated by an event bus. +ORION is intentionally structured as small services coordinated by an event bus, split across a runtime process and a client process. ```mermaid flowchart LR - U[User] --> C[CLI / TUI] - C --> O[Orchestrator] - O --> B[EventBus] - B --> S[(SQLite Event Store)] - - B --> V[Voice Recording Service] - V --> T[Transcript Generation Service] - T --> A[Agent Service] - A --> X[TTS Service] - X --> P[Audio Playback Service] - P --> D[Pipeline Complete] - - B -.global observers.-> L[Logging Service] - B -.global observers.-> UI[TUI Service] - UI --> TUI[Textual App] + subgraph client["client/ · Rust + Ratatui"] + ui["TUI
conversation · event stream · prompt"] + mic["Mic capture
cpal + hound"] + spk["Speech out
spd-say"] + end + + subgraph runtime["runtime/ · Python daemon"] + srv["IPC server
Unix socket · NDJSON"] + bridge["IPC bridge"] + stt["Transcription
Groq Whisper"] + bus["EventBus"] + store[("SQLite
event store")] + agent["AgentService"] + agentgraph["LangGraph
agent graph"] + memory["Memory module
Qdrant · Neo4j · SQLite"] + mcp["MCP manager"] + tts["TTS service
Kokoro"] + publisher["IPC publisher"] + end + + ui -->|submit_prompt| srv + mic -->|voice_end + wav path| srv + srv --> bridge + bridge --> stt + stt --> bus + bridge --> bus + bus --> store + bus --> agent + agent --> agentgraph + agentgraph <--> memory + agentgraph <--> mcp + bus --> tts + bus --> publisher + publisher -->|assistant_start / chunk / end| ui + ui --> spk ``` -### Runtime Sequence +### Agent Graph ```mermaid -sequenceDiagram - participant User - participant Orchestrator - participant Bus as EventBus - participant Store as SQLite Store - participant Voice as Voice Recording - participant STT as Transcript Generation - participant Agent as Agent Service - participant TTS as TTS Service - participant Audio as Audio Playback - - User->>Orchestrator: Start pipeline - Orchestrator->>Bus: Publish PipelineStartEvent - Bus->>Store: Append event - Bus->>Voice: Handle PipelineStartEvent - Voice->>Bus: Publish VoiceRecordingStartEvent - Voice->>Bus: Publish VoiceRecordingCompletedEvent - Bus->>STT: Handle VoiceRecordingCompletedEvent - STT->>Bus: Publish TranscriptGeneratedEvent - Bus->>Agent: Handle TranscriptGeneratedEvent - Agent->>Bus: Publish ResponseGeneratedEvent - Bus->>TTS: Handle ResponseGeneratedEvent - TTS->>Bus: Publish SpeechGeneratedEvent - Bus->>Audio: Handle SpeechGeneratedEvent - Audio->>Bus: Publish PipelineCompleteEvent +flowchart LR + s0([START]) --> retrieve + retrieve["retrieve
summary · facts · episodes · recent"] --> agent + agent["agent
LLM + bound tools"] -->|tool calls| tools + tools["tools
local automation + MCP"] --> agent + agent -->|no tool calls| remember + remember["remember
store episode · extract facts"] --> s1([END]) ``` -## Planned Platform Direction - -The codebase is designed to grow beyond voice I/O into a broader agent platform. +### Runtime Sequence ```mermaid -flowchart TB - Goal[User Goal] --> Planner[Planning / Reasoning] - Planner --> Router[Tool Router] - Router --> Memory[Memory / State] - Router --> Vision[Perception / Vision] - Router --> PC[PC Control] - Router --> Code[Code Editing / Code Generation] - Router --> Web[Browser / Web Tasks] - Router --> Audio[Voice I/O] - - PC --> Obs[Observations] - Code --> Obs - Web --> Obs - Audio --> Obs - Vision --> Obs - Obs --> Planner - Planner --> Audit[Event Log / Trace] +sequenceDiagram + participant User + participant Client as Client · Rust + participant Bridge as IPC Bridge + participant Bus as EventBus + participant Agent as AgentService + participant Graph as Agent Graph + participant Memory + + User->>Client: Type a prompt (or press v to talk) + Client->>Bridge: submit_prompt / voice_end(path) + Bridge->>Bridge: Transcribe WAV via Groq Whisper + Bridge->>Bus: ChatPipelineStartEvent + Bus->>Agent: Dispatch to subscriber + Agent->>Graph: ainvoke(state) + Graph->>Memory: retrieve(query) + Memory-->>Graph: summary · facts · episodes · recent messages + Graph->>Graph: LLM reasoning and tool execution + Graph->>Memory: remember(episode) + extracted facts + Agent->>Bus: ResponseStarted / ResponseChunk / ResponseCompleted + Bus->>Client: assistant_start / assistant_chunk / assistant_end + Client->>User: Render reply and speak it ``` -This is the direction the project is heading: - -- Computer control and desktop automation -- Code editing, patch generation, and code review workflows -- Multi-step task planning and execution -- Persistent memory and traceability -- Multi-modal perception and interaction +Every event shown above is also appended to the SQLite event store and forwarded to global observers (logging, IPC publisher). ## Repository Layout -ORION is a monorepo with two independent applications that share a single -Git repository and communicate over an IPC protocol: - -- **`runtime/`** — the Python AI runtime (daemon). -- **`client/`** — the Rust (Ratatui) terminal client. +ORION is a monorepo with two independent applications that share a single Git repository and communicate over an IPC protocol. ```text orion/ -├── runtime/ # Python AI runtime (daemon) -│ ├── src/orion/ # Installable application package -│ │ ├── __main__.py # Package entrypoint -│ │ ├── agent/ # Memory-aware agent graph and prompts -│ │ ├── bus/ # Event bus and subscription helpers -│ │ ├── cli/ # Typer CLI entrypoint -│ │ ├── core/ # Shared utilities such as the singleton metaclass -│ │ ├── events/ # Event models and registry -│ │ ├── integrations/ # External integrations such as MCP -│ │ ├── memory/ # Memory providers, planning, and persistence -│ │ ├── orchestrator/ # Runtime bootstrapping and pipeline entrypoint -│ │ ├── runtime/ # Runtime lifecycle and run loop -│ │ ├── services/ # Recording, STT, agent, TTS, playback, logging -│ │ ├── store/ # SQLite event persistence -│ │ └── transport/ # IPC protocol and bridge to the client -│ ├── tests/ # Behavior and smoke tests -│ ├── pyproject.toml # Runtime metadata and dependencies +├── runtime/ # Python AI runtime (daemon) +│ ├── src/orion/ +│ │ ├── __main__.py # Package entrypoint +│ │ ├── agent/ # Agent graph, nodes, state, prompts, local tools +│ │ │ └── nodes/ # retrieve · agent · remember · recall +│ │ ├── bus/ # Event bus and subscription helpers +│ │ ├── cli/ # Typer entrypoint (starts the daemon) +│ │ ├── core/ # Shared utilities such as the singleton metaclass +│ │ ├── events/ # Domain event models and registry +│ │ ├── integrations/_mcp/ # MCP config, server, manager, discovery, LangChain adapter +│ │ ├── llm/ # Provider abstraction (Gemini, Groq) and factory +│ │ ├── memory/ # Providers, planner, session, models +│ │ │ ├── interfaces/ # embeddings · vector · graph · summary +│ │ │ └── providers/ # sentence-transformers · qdrant · neo4j · sqlite +│ │ ├── orchestrator/ # Service wiring and startup/shutdown +│ │ ├── runtime/ # Lifecycle, runtime, application entrypoint +│ │ ├── services/ # Agent, TTS, voice recording, logging, IPC publisher +│ │ ├── store/ # SQLite event persistence +│ │ └── transport/ # IPC server, sessions, protocol, bridge, transcription +│ ├── tests/ # Behaviour and smoke tests +│ ├── mcp.json # MCP server configuration +│ ├── pyproject.toml │ └── uv.lock -├── client/ # Rust (Ratatui) terminal client -│ ├── Cargo.toml -│ └── src/main.rs -├── assets/ # Logo and visual assets -├── docker-compose.yml # Qdrant + Neo4j backends for memory -└── README.md +├── client/ # Rust (Ratatui) terminal client +│ ├── src/ +│ │ ├── main.rs # Terminal setup and async event loop +│ │ ├── app.rs # Application state and runtime-event handling +│ │ ├── ui.rs # Layout and frame composition +│ │ ├── audio.rs # Microphone capture and speech output +│ │ ├── effects.rs # tachyonfx animations +│ │ ├── theme.rs # Colours and styles +│ │ ├── ipc/ # Socket client, session, protocol, envelopes, events +│ │ └── widgets/ # header · conversation · prompt · events · status +│ └── Cargo.toml +├── assets/ # Logo and visual assets +├── docker-compose.yml # Qdrant + Neo4j backends for memory +└── .github/workflows/pytest.yml # Lint, tests, and client build ``` ## How It Works +### Runtime and Lifecycle + +`OrionRuntime` is the top-level lifecycle manager. Components are registered in order (event store → memory → MCP manager → orchestrator → IPC server), started in registration order, and shut down in reverse. A failed shutdown is collected into an `ExceptionGroup` rather than silently swallowed. + ### Event Bus -The `EventBus` is the center of the runtime. Every event is written to the store first, then fanned out to subscribed handlers and global observers. +The `EventBus` is the centre of the runtime and a process-wide singleton. Every event is written to the store first, then fanned out to handlers subscribed to that exact event type and to global observers. Handlers run concurrently via `asyncio.gather`. ### Orchestrator -The orchestrator owns startup and shutdown: +The orchestrator owns service wiring: -- builds the service list -- starts each service -- subscribes logging and UI observers -- emits a `PipelineStartEvent` -- tears everything down cleanly on exit +- builds runtime services from a `ServiceContext` (LLM, memory, MCP manager) +- registers each service's event subscriptions +- starts every service +- subscribes global observers (`LoggingService`, `IPCPublisherService`) +- tears everything down in reverse on exit ### Services -Each service is a small, focused unit: +- `AgentService` — subscribes to `ChatPipelineStartEvent`, builds the graph with per-request memory and MCP tools, and publishes the response events +- `TTSService` — subscribes to `ResponseCompletedEvent` and synthesizes speech with Kokoro into `data/audio/` +- `VoiceRecordingService` — server-side VAD recorder, retained but currently dormant (see [Current State](#current-state-and-known-gaps)) +- `LoggingService` — global observer that logs the full event stream +- `IPCPublisherService` — global observer that translates events into IPC envelopes for the originating client session + +### IPC Transport + +`IPCServer` accepts Unix socket connections and creates a `ClientSession` per client. `IPCBridge` translates in both directions: inbound envelopes become domain events (`submit_prompt` → `ChatPipelineStartEvent`, `voice_end` → transcribe then `ChatPipelineStartEvent`), and outbound events become envelopes routed to the right session. Messages are NDJSON-encoded Pydantic models, so both ends validate what they receive. + +### Memory + +`MemoryModule` owns four providers behind interfaces: + +| Concern | Provider | Backend | +| --- | --- | --- | +| Embeddings | `SentenceTransformerEmbeddingProvider` | local model (default `BAAI/bge-small-en-v1.5`) | +| Semantic recall | `QdrantVectorMemory` | Qdrant | +| Facts | `Neo4jKnowledgeGraph` | Neo4j | +| Rolling summary | `SQLiteSummaryStore` | SQLite | + +A `MemorySession` is bound to one pipeline execution. `retrieve()` fans out across the summary, graph facts, semantic matches, and recent history in parallel, and publishes an event for each step so the client can trace it. After each turn, `RememberNode` stores the episode and asks the LLM to extract stable long-term facts about the user, which are written to the graph. + +`RetrievalPlanner` is wired up on both the memory module and the agent service but is not called yet — retrieval currently uses the fixed fan-out above rather than an LLM-chosen strategy. + +### MCP Integrations + +`MCPManager` starts every enabled server from `mcp.json`, discovers its tools, keeps a tool → server routing table (first server to claim a name wins), and exposes the tools to the agent as LangChain tools. A server that fails to start does not prevent the others from starting. Every server and tool-call outcome is published as an event. + +### Client -- `VoiceRecordingService` records microphone input and writes `data/audio/input.wav` -- `TranscriptGenerationService` sends audio to Groq Whisper and emits text -- `AgentService` turns the transcript into a response -- `TTSService` synthesizes speech into `data/audio/output.wav` -- `AudioPlaybackService` plays the response and closes the pipeline -- `TUIService` mirrors the live stream into the terminal UI -- `LoggingService` is available as a global observer hook +The client runs a single-threaded Tokio runtime (the audio input stream is `!Send`) and drives a `tokio::select!` loop over three sources: a ~30 FPS render tick, terminal input, and the IPC event stream. State lives in `App`, rendering lives in `ui.rs`, and `effects.rs` post-processes the composed buffer. ## Requirements -- Python 3.14+ -- Audio input device -- Audio output device -- `GROQ_API_KEY` configured in the environment +- Python 3.11+ and [`uv`](https://docs.astral.sh/uv/) +- Rust 1.85+ (the client uses edition 2024) +- Docker (for the Qdrant and Neo4j memory backends) +- An API key for the configured LLM provider — `GEMINI_API_KEY` or `GROQ_API_KEY` +- `GROQ_API_KEY` for voice input (transcription always uses Groq Whisper) +- Audio input device (voice input) and `speech-dispatcher` / `spd-say` (spoken replies) — both optional; typing works without either +- On Linux, ALSA headers for the client's audio capture: `libasound2-dev` (Debian/Ubuntu) or `alsa-lib` (Arch) +- Node.js / `npx` if you keep the default filesystem MCP server -## Installation +The first runtime start downloads the embedding model and Kokoro voice weights, so expect it to take a while. -The Python runtime lives in `runtime/`: +## Installation ```bash +# Python runtime cd runtime uv sync + +# Rust client +cd ../client +cargo build ``` If you are not using `uv`, install the dependencies from `runtime/pyproject.toml` using your preferred Python tooling. ## Configuration -Create a `.env` file inside `runtime/` with your API key: +Create a `.env` file inside `runtime/`: ```env -GROQ_API_KEY=your_key_here +# LLM +ORION_LLM_PROVIDER=gemini # gemini | groq +ORION_LLM_MODEL=gemini-3.1-flash-lite # required — no default +GEMINI_API_KEY=your_key_here +GROQ_API_KEY=your_key_here # also used for voice transcription + +# Memory +NEO4J_PASSWORD=orion123 # matches docker-compose.yml ``` -Optional local artifacts: +Everything else has a working local default: + +| Variable | Default | Purpose | +| --- | --- | --- | +| `ORION_LLM_PROVIDER` | `gemini` | Which provider the LLM factory builds | +| `ORION_LLM_MODEL` | — | Model id; the provider raises if unset | +| `GEMINI_API_KEY` | — | Required when the provider is `gemini` | +| `GROQ_API_KEY` | — | Required when the provider is `groq`, and for transcription | +| `ORION_SQLITE_DB` | `orion.db` | SQLite database for the rolling summary | +| `ORION_EMBEDDING_MODEL` | `BAAI/bge-small-en-v1.5` | Sentence-transformer embedding model | +| `QDRANT_HOST` / `QDRANT_PORT` | `localhost` / `6333` | Qdrant connection | +| `NEO4J_URI` | `bolt://localhost:7687` | Neo4j connection | +| `NEO4J_USERNAME` / `NEO4J_PASSWORD` | `neo4j` / empty | Neo4j credentials | +| `MCP_CONFIG` | `mcp.json` | Path to the MCP server configuration | +| `SOCKET_PATH` | `/tmp/orion.sock` | Unix socket the runtime listens on | + +> The client's socket path is currently compiled in as `/tmp/orion.sock`, so changing `SOCKET_PATH` also requires changing `SOCKET_PATH` in `client/src/main.rs`. + +### MCP Servers + +`runtime/mcp.json` configures the MCP servers to start. `${PROJECT_ROOT}` resolves to the directory containing the config file, and environment variables are expanded: + +```json +{ + "mcpServers": { + "filesystem": { + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "${PROJECT_ROOT}"], + "enabled": true + } + } +} +``` -- `orion.db` - SQLite event store -- `data/audio/input.wav` - captured microphone input -- `data/audio/output.wav` - generated response audio +`stdio` (via `command` + `args`) and `http` (via `"type": "http"` + `url`) transports are supported. Set `"enabled": false` to skip a server. + +### Local Artifacts + +- `orion.db` — SQLite event store and summary memory +- `data/audio/*.wav` — synthesized speech output +- `logs/orion.log` — runtime log +- `$TMPDIR/orion_client_recording.wav` — the client's most recent recording ## Run -Start the memory backends (from the repo root), then run the runtime: +Start the memory backends, then the runtime, then the client — the runtime and client need separate terminals. ```bash +# 1. Memory backends (repo root) docker compose up -d # Qdrant + Neo4j -cd runtime -uv run python -m orion # or: uv run orion -``` - -## Client (Rust) -The terminal client lives in `client/` and is built with Cargo. Start the -runtime first (it opens the IPC socket the client connects to), then: +# 2. Runtime daemon +cd runtime +uv run orion # or: uv run python -m orion -```bash +# 3. Terminal client (second terminal) cd client cargo run ``` -The client is a Ratatui TUI that talks to the runtime over `/tmp/orion.sock`. -It renders the conversation, a live event stream, and Copilot-style activity -logs, with tachyonfx animations. - -### Keybindings - -| Key | Action | -|-----|--------| -| `i` | enter insert mode (type a prompt) | -| `Enter` | send the typed prompt (insert mode) | -| `Esc` | back to normal mode | -| `v` | push-to-talk: press to start recording, press again to send | -| `s` | stop the assistant's speech | -| `j`/`k`, arrows, `PgUp`/`PgDn` | scroll the conversation | -| `q` | quit | +The runtime prints a startup banner and waits for connections. If the client starts first, or the socket is missing, it shows `OFFLINE` in the status bar and does not retry — start the runtime and relaunch the client. -### Voice +## Client Keybindings -Voice is handled entirely by the client; the runtime never touches audio -hardware: +The client is modal, like vim. The status bar shows the current mode, connection state, activity, and event count. -1. Press `v` to record from your microphone, `v` again to stop. -2. The client saves a WAV and sends its path to the runtime over IPC. -3. The runtime transcribes it (Groq Whisper) and runs the normal chat pipeline. -4. The response streams back and the client speaks it aloud. +**Normal mode** -Text-to-speech uses your system speech engine via **speech-dispatcher**. Install -it to hear responses (otherwise TTS is silently skipped): - -```bash -sudo pacman -S speech-dispatcher # Arch / EndeavourOS -# Debian/Ubuntu: sudo apt-get install -y speech-dispatcher -``` +| Key | Action | +| --- | --- | +| `i` / `a` | Enter insert mode | +| `v` | Toggle voice recording (press again to send) | +| `s` | Interrupt assistant speech | +| `k` / `↑`, `j` / `↓` | Scroll the conversation | +| `PageUp` / `PageDown` | Scroll by 5 lines | +| `Home` | Jump to the top | +| `End` / `G` | Jump to the bottom | +| `q` / `Esc` | Quit | + +**Insert mode** -Building the client also needs the ALSA development headers for microphone -capture (`cpal`): +| Key | Action | +| --- | --- | +| `Enter` | Send the prompt | +| `Esc` | Back to normal mode | +| `←` / `→` | Move the cursor | +| `Backspace` | Delete a character | -```bash -sudo pacman -S alsa-lib # Arch / EndeavourOS -# Debian/Ubuntu: sudo apt-get install -y libasound2-dev -``` +`Ctrl+C` quits from either mode. Clicking the prompt box enters insert mode; clicking elsewhere returns to normal mode. The scroll wheel scrolls whichever panel is under the pointer. ## Tests ```bash cd runtime uv run pytest + +cd ../client +cargo test ``` ## Before Pushing a PR -Run the project checks locally before you push a pull request: +CI runs ruff, pytest, and a client build on every push and pull request. Run the same checks locally first: ```bash cd runtime uv run pytest uv run ruff check . uv run ruff format --check . + +cd ../client +cargo build +cargo test ``` -If you are making code changes, it is worth running both commands again after your final edit so the PR starts clean. +If you are making code changes, it is worth running these again after your final edit so the PR starts clean. + +## Current State and Known Gaps + +The runtime is further along than the wiring between the two processes, so a few things are deliberately half-connected: + +- **Responses are not truly token-streamed yet.** The agent publishes one `assistant_chunk` containing the whole response; the protocol and client already handle incremental chunks. +- **Kokoro output does not reach the client.** `TTSService` still synthesizes WAVs, but the speech events carry no IPC message type, so nothing is forwarded. The client speaks replies itself with `spd-say`. +- **The client only decodes a subset of messages.** `assistant_start`/`chunk`/`end`, `status`, and `error` map to typed events; everything else falls through to `Unknown`. The handlers for tool, pipeline, and voice events exist but are not yet reachable end to end. +- **Server-side voice recording is dormant.** `VoiceRecordingService` subscribes to the bare `PipelineStartEvent`, which nothing publishes anymore, and `TranscriptGenerationService` is commented out of the service list. The client owns the microphone and the runtime transcribes the file it is handed. +- **Sessions are per-connection.** Memory is global rather than scoped per client, and there is no reconnect or cancellation handling yet (`cancel_request` is defined in the protocol but not implemented). +- **`execute_shell_command` runs unsandboxed shell commands** on the host as part of the agent's tool set. Treat the runtime as trusted-local-use only. + +## Planned Platform Direction + +The codebase is designed to grow beyond the current loop into a broader agent platform. -## Notes +```mermaid +flowchart TB + Goal[User Goal] --> Planner[Planning / Reasoning] + Planner --> Router[Tool Router] + Router --> Memory[Memory / State] + Router --> Vision[Perception / Vision] + Router --> PC[PC Control] + Router --> Code[Code Editing / Code Generation] + Router --> Web[Browser / Web Tasks] + Router --> Audio[Voice I/O] -- `chat` is still a placeholder command. -- `doctor` is intentionally minimal right now. -- The current voice loop is continuous; it keeps starting new pipelines until the TUI exits. -- The codebase is already structured for more advanced agents, but desktop control and coding automation are still future work. + PC --> Obs[Observations] + Code --> Obs + Web --> Obs + Audio --> Obs + Vision --> Obs + Obs --> Planner + Planner --> Audit[Event Log / Trace] +``` + +This is the direction the project is heading: + +- Computer control and desktop automation +- Code editing, patch generation, and code review workflows +- Multi-step task planning and execution +- Real token streaming and cancellable requests +- Multi-modal perception and interaction ## Vision diff --git a/client/src/app.rs b/client/src/app.rs index 02940cf..60268c9 100644 --- a/client/src/app.rs +++ b/client/src/app.rs @@ -228,7 +228,12 @@ impl App { pub fn handle_runtime_event(&mut self, event: RuntimeEvent) { match event { - RuntimeEvent::Connected => self.on_connected(), + // ------------------------------------------------------------------ + // Connection + // ------------------------------------------------------------------ + RuntimeEvent::Connected => { + self.on_connected(); + } RuntimeEvent::Disconnected => { self.mode = "DISCONNECTED".into(); @@ -237,60 +242,206 @@ impl App { .push("DISCONNECTED", EventStatus::Failed, "runtime disconnected"); } - RuntimeEvent::UserPrompt(text) => { - let text = text.trim().to_string(); - // Ignore the echo of a prompt we already showed locally (typed); - // display it when it's new — i.e. a transcribed voice message. - if !text.is_empty() - && self.conversation.last_user_text().as_deref() != Some(text.as_str()) - { - self.msg_counter += 1; - self.conversation.add_message(Message::new( - format!("msg-{}", self.msg_counter), - Author::User, - text, - )); - self.effects.on_message(); - } + // ------------------------------------------------------------------ + // Pipeline + // ------------------------------------------------------------------ + RuntimeEvent::PipelineStart => { + self.mode = "PROCESSING".into(); + self.events + .push("PIPELINE", EventStatus::Running, "pipeline started"); + } + + RuntimeEvent::VoicePipelineStart => { + self.mode = "VOICE PIPELINE".into(); + self.events.push( + "VOICE_PIPELINE", + EventStatus::Running, + "voice pipeline started", + ); + } + + RuntimeEvent::ChatPipelineStart => { + self.mode = "THINKING".into(); + self.events.push( + "CHAT_PIPELINE", + EventStatus::Running, + "chat pipeline started", + ); + } + + RuntimeEvent::PipelineComplete => { + self.mode = "IDLE".into(); + self.events + .bump_last("PIPELINE", EventStatus::Completed, "pipeline complete"); + } + + RuntimeEvent::PipelineFailed { error } => { + self.mode = "PIPELINE ERROR".into(); + self.effects.on_status_change(theme::DANGER); + + self.events + .push("PIPELINE_FAILED", EventStatus::Failed, error.clone()); + + self.add_activity(format!("Pipeline failed: {}", error), ActivityKind::Failed); + } + + RuntimeEvent::PipelineRestart => { + self.mode = "RESTARTING".into(); + self.events.push( + "PIPELINE_RESTART", + EventStatus::Running, + "pipeline restarting", + ); + } + + // ------------------------------------------------------------------ + // Voice + // ------------------------------------------------------------------ + RuntimeEvent::VoiceRecordingStart => { + self.mode = "RECORDING".into(); + self.effects.on_status_change(theme::DANGER); + + self.events + .push("VOICE_RECORDING", EventStatus::Running, "recording started"); } + RuntimeEvent::VoiceRecordingCompleted { audio_path } => { + self.mode = "TRANSCRIBING".into(); + + self.events.push( + "VOICE_RECORDING", + EventStatus::Completed, + audio_path.unwrap_or_else(|| "recording completed".into()), + ); + } + + RuntimeEvent::VoiceRecordingFailed { error } => { + self.mode = "REC ERROR".into(); + self.effects.on_status_change(theme::DANGER); + + self.events + .push("VOICE_RECORDING", EventStatus::Failed, error.clone()); + + self.add_activity(format!("Recording failed: {}", error), ActivityKind::Failed); + } + + RuntimeEvent::SpeechDetected => { + self.events + .push("SPEECH", EventStatus::Running, "speech detected"); + } + + RuntimeEvent::SilenceDetected { silence_duration } => { + self.events.push( + "SILENCE", + EventStatus::Info, + format!("silence detected ({:.2}s)", silence_duration), + ); + } + + // ------------------------------------------------------------------ + // Speech-to-Text + // ------------------------------------------------------------------ + RuntimeEvent::TranscriptGenerated { text } => { + self.mode = "THINKING".into(); + + self.events + .push("TRANSCRIPT", EventStatus::Completed, text.clone()); + + self.conversation.add_message(Message::new( + format!("msg-{}", self.msg_counter + 1), + Author::User, + text, + )); + + self.msg_counter += 1; + self.effects.on_message(); + } + + RuntimeEvent::TranscriptGenerationFailed { error } => { + self.mode = "TRANSCRIPTION ERROR".into(); + self.effects.on_status_change(theme::DANGER); + + self.events + .push("TRANSCRIPT", EventStatus::Failed, error.clone()); + + self.add_activity( + format!("Transcription failed: {}", error), + ActivityKind::Failed, + ); + } + + // ------------------------------------------------------------------ + // Agent + // ------------------------------------------------------------------ + RuntimeEvent::AgentProcessingStart => { + self.mode = "THINKING".into(); + + self.events + .push("AGENT", EventStatus::Running, "agent processing started"); + } + + // ------------------------------------------------------------------ + // Assistant + // ------------------------------------------------------------------ RuntimeEvent::AssistantStart => { self.mode = "RESPONDING".into(); self.msg_counter += 1; - // Begin a single message bubble for the streaming response self.conversation .begin_assistant_message(format!("msg-{}", self.msg_counter)); + self.effects.on_message(); + self.events .push("RESPONSE", EventStatus::Running, "assistant responding"); } RuntimeEvent::AssistantChunk(text) => { self.conversation.append_assistant_chunk(&text); - // Coalesce chunk spam onto the running RESPONSE trace line. + self.events .bump_last("RESPONSE", EventStatus::Running, "streaming…"); } RuntimeEvent::AssistantEnd => { self.mode = "IDLE".into(); + self.conversation.finish_assistant_message(); + self.events .bump_last("RESPONSE", EventStatus::Completed, "response complete"); - // Client-side TTS: speak the completed response. + self.speaker.speak(&self.conversation.last_assistant_text()); } + RuntimeEvent::ResponseGenerationFailed { error } => { + self.mode = "RESPONSE ERROR".into(); + self.effects.on_status_change(theme::DANGER); + + self.events + .push("RESPONSE", EventStatus::Failed, error.clone()); + + self.add_activity( + format!("Response generation failed: {}", error), + ActivityKind::Failed, + ); + } + + // ------------------------------------------------------------------ + // Tools + // ------------------------------------------------------------------ RuntimeEvent::ToolStarted { name } => { self.mode = format!("TOOL: {}", name); + self.events .push("TOOL_STARTED", EventStatus::Running, name.clone()); + self.add_activity(format!("Running {}…", name), ActivityKind::Running); } RuntimeEvent::ToolFinished { name, success } => { let label = if success { "OK" } else { "FAILED" }; + self.mode = format!("TOOL {}: {}", label, name); let status = if success { @@ -298,6 +449,7 @@ impl App { } else { EventStatus::Failed }; + self.events.push("TOOL_FINISHED", status, name.clone()); if success { @@ -307,39 +459,124 @@ impl App { } } + // ------------------------------------------------------------------ + // Text-to-Speech + // ------------------------------------------------------------------ + RuntimeEvent::SpeechSynthesisStart { text } => { + self.mode = "SYNTHESIZING".into(); + + self.events.push("TTS", EventStatus::Running, text); + } + + RuntimeEvent::SpeechGenerated { audio_path, text } => { + self.mode = "IDLE".into(); + + self.events + .push("TTS", EventStatus::Completed, audio_path.unwrap_or(text)); + } + + RuntimeEvent::SpeechGenerationFailed { error } => { + self.mode = "TTS ERROR".into(); + self.effects.on_status_change(theme::DANGER); + + self.events.push("TTS", EventStatus::Failed, error.clone()); + + self.add_activity( + format!("Speech synthesis failed: {}", error), + ActivityKind::Failed, + ); + } + + // ------------------------------------------------------------------ + // Audio Playback + // ------------------------------------------------------------------ + RuntimeEvent::AudioPlaybackStarted => { + self.mode = "PLAYING".into(); + + self.events + .push("PLAYBACK", EventStatus::Running, "audio playback started"); + } + + RuntimeEvent::AudioPlaybackCompleted => { + self.mode = "IDLE".into(); + + self.events.bump_last( + "PLAYBACK", + EventStatus::Completed, + "audio playback complete", + ); + } + + RuntimeEvent::AudioPlaybackFailed { error } => { + self.mode = "PLAYBACK ERROR".into(); + self.effects.on_status_change(theme::DANGER); + + self.events + .push("PLAYBACK", EventStatus::Failed, error.clone()); + + self.add_activity( + format!("Audio playback failed: {}", error), + ActivityKind::Failed, + ); + } + + // ------------------------------------------------------------------ + // Runtime + // ------------------------------------------------------------------ RuntimeEvent::Status(status) => { self.mode = status.clone(); + self.events.push("STATUS", EventStatus::Info, status); } RuntimeEvent::Error { code, message } => { self.mode = format!("ERROR: {}", message); self.effects.on_status_change(theme::DANGER); + self.events.push( "ERROR", EventStatus::Failed, format!("{}: {}", code, message), ); + + self.add_activity(format!("{}: {}", code, message), ActivityKind::Failed); } - // Heartbeat — intentionally not traced (too noisy). - RuntimeEvent::Ping | RuntimeEvent::Pong => {} + // ------------------------------------------------------------------ + // Heartbeat + // ------------------------------------------------------------------ + RuntimeEvent::Ping | RuntimeEvent::Pong => { + // Intentionally not traced. + } + // ------------------------------------------------------------------ + // Voice streaming + // ------------------------------------------------------------------ RuntimeEvent::VoiceStart => { self.mode = "VOICE RECORDING".into(); + self.events .push("VOICE_START", EventStatus::Running, "recording"); } - RuntimeEvent::VoiceChunk { .. } => {} + RuntimeEvent::VoiceChunk { .. } => { + // Audio chunks are intentionally not logged individually. + } RuntimeEvent::VoiceEnd => { self.mode = "PROCESSING VOICE".into(); + self.events .push("VOICE_END", EventStatus::Info, "processing"); } - RuntimeEvent::Unknown(_) => {} + // ------------------------------------------------------------------ + // Unknown + // ------------------------------------------------------------------ + RuntimeEvent::Unknown(_) => { + self.events + .push("UNKNOWN", EventStatus::Info, "unsupported runtime event"); + } } } diff --git a/client/src/ipc/events.rs b/client/src/ipc/events.rs index 28168da..8023ccc 100644 --- a/client/src/ipc/events.rs +++ b/client/src/ipc/events.rs @@ -6,16 +6,62 @@ pub enum RuntimeEvent { Connected, Disconnected, - /// The user's turn text echoed by the runtime — a typed prompt or a - /// transcribed voice message. - UserPrompt(String), + // ------------------------------------------------------------------ + // Pipeline + // ------------------------------------------------------------------ + PipelineStart, + VoicePipelineStart, + ChatPipelineStart, + + PipelineComplete, + PipelineFailed { + error: String, + }, + PipelineRestart, + + // ------------------------------------------------------------------ + // Voice + // ------------------------------------------------------------------ + VoiceRecordingStart, + VoiceRecordingCompleted { + audio_path: Option, + }, + VoiceRecordingFailed { + error: String, + }, + + SpeechDetected, + SilenceDetected { + silence_duration: f32, + }, + + // ------------------------------------------------------------------ + // Speech-to-Text + // ------------------------------------------------------------------ + TranscriptGenerated { + text: String, + }, + TranscriptGenerationFailed { + error: String, + }, + + // ------------------------------------------------------------------ + // Agent + // ------------------------------------------------------------------ + AgentProcessingStart, // Assistant AssistantStart, AssistantChunk(String), AssistantEnd, + ResponseGenerationFailed { + error: String, + }, + + // ------------------------------------------------------------------ // Tools + // ------------------------------------------------------------------ ToolStarted { name: String, }, @@ -24,18 +70,47 @@ pub enum RuntimeEvent { success: bool, }, + // ------------------------------------------------------------------ + // Text-to-Speech + // ------------------------------------------------------------------ + SpeechSynthesisStart { + text: String, + }, + SpeechGenerated { + audio_path: Option, + text: String, + }, + SpeechGenerationFailed { + error: String, + }, + + // ------------------------------------------------------------------ + // Audio Playback + // ------------------------------------------------------------------ + AudioPlaybackStarted, + AudioPlaybackCompleted, + AudioPlaybackFailed { + error: String, + }, + + // ------------------------------------------------------------------ // Runtime + // ------------------------------------------------------------------ Status(String), Error { code: String, message: String, }, + // ------------------------------------------------------------------ // Connection + // ------------------------------------------------------------------ Ping, Pong, - // Voice (future) + // ------------------------------------------------------------------ + // Voice streaming + // ------------------------------------------------------------------ VoiceStart, VoiceChunk { sequence: u64, @@ -43,6 +118,8 @@ pub enum RuntimeEvent { }, VoiceEnd, + // ------------------------------------------------------------------ // Unknown / unsupported + // ------------------------------------------------------------------ Unknown(Envelope), } diff --git a/client/src/ipc/messages.rs b/client/src/ipc/messages.rs index f2b9f99..80604d3 100644 --- a/client/src/ipc/messages.rs +++ b/client/src/ipc/messages.rs @@ -6,48 +6,75 @@ use uuid::Uuid; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum MessageType { + // ---------------------------------------------------------------------- // Connection + // ---------------------------------------------------------------------- + Ping, Pong, + // ---------------------------------------------------------------------- // Prompt submission + // ---------------------------------------------------------------------- + SubmitPrompt, CancelRequest, + // ---------------------------------------------------------------------- // Voice + // ---------------------------------------------------------------------- + VoiceStart, VoiceChunk, VoiceEnd, + // ---------------------------------------------------------------------- // Assistant streaming + // ---------------------------------------------------------------------- + AssistantStart, AssistantChunk, AssistantEnd, + // ---------------------------------------------------------------------- // Tool execution + // ---------------------------------------------------------------------- + ToolStarted, ToolFinished, + // ---------------------------------------------------------------------- // Runtime + // ---------------------------------------------------------------------- + Status, Error, } -/// Every IPC message is wrapped inside an Envelope. +/// Every IPC message exchanged between Orion and a client is wrapped inside +/// an Envelope. +/// +/// The envelope contains transport-level metadata while the payload contains +/// the message-specific data. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Envelope { + /// IPC protocol version. #[serde(default = "default_version")] pub version: u32, + /// Unique identifier for this individual message. #[serde(default = "Uuid::new_v4")] pub id: Uuid, + /// Identifier connecting related messages to the same request/pipeline. #[serde(default = "Uuid::new_v4")] pub correlation_id: Uuid, + /// Type of message contained in the envelope. #[serde(rename = "type")] pub message_type: MessageType, + /// Message-specific payload. #[serde(default)] pub payload: Value, } @@ -57,89 +84,244 @@ const fn default_version() -> u32 { } impl Envelope { - /// Create a ping request. - pub fn ping() -> Self { + /// Create an envelope with a newly generated message ID and correlation ID. + pub fn new(message_type: MessageType, payload: Value) -> Self { Self { version: 1, id: Uuid::new_v4(), correlation_id: Uuid::new_v4(), - message_type: MessageType::Ping, - payload: json!({}), + message_type, + payload, } } - /// Create a pong response. - pub fn pong() -> Self { + /// Create an envelope using an existing correlation ID. + /// + /// This is important for runtime events because all events belonging to + /// one pipeline must remain associated with the same request. + pub fn with_correlation_id( + message_type: MessageType, + correlation_id: Uuid, + payload: Value, + ) -> Self { Self { version: 1, id: Uuid::new_v4(), - correlation_id: Uuid::new_v4(), - message_type: MessageType::Pong, - payload: json!({}), + correlation_id, + message_type, + payload, } } + /// Create a ping request. + pub fn ping() -> Self { + Self::new(MessageType::Ping, json!({})) + } + + /// Create a pong response. + pub fn pong() -> Self { + Self::new(MessageType::Pong, json!({})) + } + /// Create a prompt submission. pub fn submit_prompt(text: impl Into) -> Self { - Self { - version: 1, - id: Uuid::new_v4(), - correlation_id: Uuid::new_v4(), - message_type: MessageType::SubmitPrompt, - payload: serde_json::to_value(SubmitPromptPayload { text: text.into() }).unwrap(), - } + Self::new( + MessageType::SubmitPrompt, + serde_json::to_value(SubmitPromptPayload { + text: text.into(), + }) + .expect("failed to serialize SubmitPromptPayload"), + ) } - /// Announce the start of a voice recording session (metadata only). + /// Create a prompt submission associated with an existing request. + pub fn submit_prompt_with_correlation( + correlation_id: Uuid, + text: impl Into, + ) -> Self { + Self::with_correlation_id( + MessageType::SubmitPrompt, + correlation_id, + serde_json::to_value(SubmitPromptPayload { + text: text.into(), + }) + .expect("failed to serialize SubmitPromptPayload"), + ) + } + + /// Create a cancel request. + pub fn cancel_request(request_id: Uuid) -> Self { + Self::new( + MessageType::CancelRequest, + serde_json::to_value(CancelRequestPayload { request_id }) + .expect("failed to serialize CancelRequestPayload"), + ) + } + + /// Create a voice recording start message. pub fn voice_start(sample_rate: u32, channels: u32) -> Self { - Self { - version: 1, - id: Uuid::new_v4(), - correlation_id: Uuid::new_v4(), - message_type: MessageType::VoiceStart, - payload: serde_json::to_value(VoiceStartPayload { + Self::new( + MessageType::VoiceStart, + serde_json::to_value(VoiceStartPayload { sample_rate, channels, encoding: "pcm16".to_string(), }) - .unwrap(), - } + .expect("failed to serialize VoiceStartPayload"), + ) + } + + /// Create a voice chunk message. + pub fn voice_chunk(sequence: u64, data: Vec) -> Self { + Self::new( + MessageType::VoiceChunk, + serde_json::to_value(VoiceChunkPayload { sequence, data }) + .expect("failed to serialize VoiceChunkPayload"), + ) } - /// Finish a recording by handing the runtime the recorded file path. + /// Finish a voice recording by providing the recorded file path. pub fn voice_end(path: impl Into) -> Self { - Self { - version: 1, - id: Uuid::new_v4(), - correlation_id: Uuid::new_v4(), - message_type: MessageType::VoiceEnd, - payload: serde_json::to_value(VoiceEndPayload { path: path.into() }).unwrap(), - } + Self::new( + MessageType::VoiceEnd, + serde_json::to_value(VoiceEndPayload { path: path.into() }) + .expect("failed to serialize VoiceEndPayload"), + ) + } + + /// Signal that assistant response generation has started. + pub fn assistant_start(correlation_id: Uuid) -> Self { + Self::with_correlation_id( + MessageType::AssistantStart, + correlation_id, + json!({}), + ) + } + + /// Send an assistant response chunk. + pub fn assistant_chunk( + correlation_id: Uuid, + text: impl Into, + ) -> Self { + Self::with_correlation_id( + MessageType::AssistantChunk, + correlation_id, + serde_json::to_value(AssistantChunkPayload { + text: text.into(), + }) + .expect("failed to serialize AssistantChunkPayload"), + ) + } + + /// Signal that assistant response generation has completed. + pub fn assistant_end(correlation_id: Uuid) -> Self { + Self::with_correlation_id( + MessageType::AssistantEnd, + correlation_id, + json!({}), + ) } - /// Deserialize the payload into a strongly typed struct. + /// Signal that a tool has started. + pub fn tool_started( + correlation_id: Uuid, + name: impl Into, + ) -> Self { + Self::with_correlation_id( + MessageType::ToolStarted, + correlation_id, + serde_json::to_value(ToolStartedPayload { + name: name.into(), + }) + .expect("failed to serialize ToolStartedPayload"), + ) + } + + /// Signal that a tool has finished. + pub fn tool_finished( + correlation_id: Uuid, + name: impl Into, + success: bool, + ) -> Self { + Self::with_correlation_id( + MessageType::ToolFinished, + correlation_id, + serde_json::to_value(ToolFinishedPayload { + name: name.into(), + success, + }) + .expect("failed to serialize ToolFinishedPayload"), + ) + } + + /// Send a runtime status update. + pub fn status( + correlation_id: Uuid, + message: impl Into, + ) -> Self { + Self::with_correlation_id( + MessageType::Status, + correlation_id, + serde_json::to_value(StatusPayload { + message: message.into(), + }) + .expect("failed to serialize StatusPayload"), + ) + } + + /// Send an error to the client. + pub fn error( + correlation_id: Uuid, + code: impl Into, + message: impl Into, + ) -> Self { + Self::with_correlation_id( + MessageType::Error, + correlation_id, + serde_json::to_value(ErrorPayload { + code: code.into(), + message: message.into(), + }) + .expect("failed to serialize ErrorPayload"), + ) + } + + /// Deserialize the envelope payload into a strongly typed payload. pub fn payload(&self) -> serde_json::Result where T: DeserializeOwned, { serde_json::from_value(self.payload.clone()) } + + /// Serialize this envelope to JSON. + pub fn to_json(&self) -> serde_json::Result { + serde_json::to_string(self) + } + + /// Deserialize an envelope from JSON. + pub fn from_json(json: &str) -> serde_json::Result { + serde_json::from_str(json) + } } /* -------------------------------------------------------------------------- */ /* Payload Definitions */ /* -------------------------------------------------------------------------- */ +/// User prompt submitted to Orion. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SubmitPromptPayload { pub text: String, } +/// Request cancellation. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CancelRequestPayload { pub request_id: Uuid, } +/// Metadata describing a voice recording session. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct VoiceStartPayload { pub sample_rate: u32, @@ -147,53 +329,65 @@ pub struct VoiceStartPayload { pub encoding: String, } +/// A streamed audio frame. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct VoiceChunkPayload { pub sequence: u64, pub data: Vec, } +/// Marks the end of a voice recording. +/// +/// The runtime reads and transcribes the specified file. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct VoiceEndPayload { - /// Path to the recorded audio file the runtime should transcribe. pub path: String, } +/// Signals the start of assistant response generation. #[derive(Debug, Clone, Serialize, Deserialize, Default)] -pub struct AssistantStartPayload; +pub struct AssistantStartPayload {} +/// A streamed fragment of assistant output. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AssistantChunkPayload { pub text: String, } +/// Signals completion of assistant response generation. #[derive(Debug, Clone, Serialize, Deserialize, Default)] -pub struct AssistantEndPayload; +pub struct AssistantEndPayload {} +/// Indicates that a tool has started executing. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolStartedPayload { pub name: String, } +/// Indicates that a tool has completed. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolFinishedPayload { pub name: String, pub success: bool, } +/// General runtime status update. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct StatusPayload { pub message: String, } +/// Runtime error returned to the client. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ErrorPayload { pub code: String, pub message: String, } +/// Ping payload. #[derive(Debug, Clone, Serialize, Deserialize, Default)] -pub struct PingPayload; +pub struct PingPayload {} +/// Pong payload. #[derive(Debug, Clone, Serialize, Deserialize, Default)] -pub struct PongPayload; +pub struct PongPayload {} \ No newline at end of file diff --git a/runtime/mcp.json b/runtime/mcp.json index 4fdd261..d8c1f48 100644 --- a/runtime/mcp.json +++ b/runtime/mcp.json @@ -7,17 +7,7 @@ "@modelcontextprotocol/server-filesystem", "${PROJECT_ROOT}" ], - "enabled": true, - "tools": [ - "read_text_file", - "list_directory", - "directory_tree", - "search_files", - "write_file", - "edit_file", - "move_file", - "create_directory" - ] + "enabled": true } } } diff --git a/runtime/pyproject.toml b/runtime/pyproject.toml index e7311fc..a2addc2 100644 --- a/runtime/pyproject.toml +++ b/runtime/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "kokoro>=0.9.4", "langchain>=1.3.11", "langchain-core>=1.4.8", + "langchain-google-genai>=4.3.5", "langchain-groq>=1.1.3", "langchain-mcp-adapters>=0.3.0", "langgraph>=1.2.8", diff --git a/runtime/src/orion/agent/nodes/remember.py b/runtime/src/orion/agent/nodes/remember.py index 83a6929..3f7436c 100644 --- a/runtime/src/orion/agent/nodes/remember.py +++ b/runtime/src/orion/agent/nodes/remember.py @@ -10,6 +10,7 @@ from orion.memory.models import ConversationEpisode, Fact from orion.memory.session import MemorySession + MEMORY_PROMPT = """ You are an information extraction system. @@ -64,14 +65,48 @@ """ +def message_content_to_text(content: object) -> str: + """ + Convert LangChain message content into plain text. + + Chat models may return: + - a plain string + - a list of content blocks + - dictionaries containing text + """ + + if isinstance(content, str): + return content + + if not isinstance(content, list): + return "" + + parts: list[str] = [] + + for block in content: + if isinstance(block, str): + parts.append(block) + continue + + if not isinstance(block, dict): + continue + + text = block.get("text") + + if isinstance(text, str): + parts.append(text) + + return "".join(parts) + + class RememberNode: def __init__( self, memory: MemorySession, llm: BaseChatModel, ) -> None: - self.memory: MemorySession = memory - self.llm: BaseChatModel = llm + self.memory = memory + self.llm = llm async def __call__( self, @@ -81,25 +116,43 @@ async def __call__( messages = state["messages"] user = next( - message - for message in reversed(messages) - if isinstance(message, HumanMessage) + ( + message + for message in reversed(messages) + if isinstance(message, HumanMessage) + ), + None, ) assistant = next( - message for message in reversed(messages) if isinstance(message, AIMessage) + ( + message + for message in reversed(messages) + if isinstance(message, AIMessage) + ), + None, ) - assert isinstance(user.content, str) - assert isinstance(assistant.content, str) + if user is None or assistant is None: + return {} + + user_content = message_content_to_text( + user.content, + ) + + assistant_content = message_content_to_text( + assistant.content, + ) + + if not user_content or not assistant_content: + return {} episode = ConversationEpisode( correlation_id=state["correlation_id"], - user_message=user.content, - assistant_message=assistant.content, + user_message=user_content, + assistant_message=assistant_content, ) - # Store semantic conversation memory await self.memory.remember(episode) response = await self.llm.ainvoke( @@ -107,35 +160,70 @@ async def __call__( ("system", MEMORY_PROMPT), ( "human", - f"User: {user.content}\nAssistant: {assistant.content}", + f"User: {user_content}\n" + f"Assistant: {assistant_content}", ), ] ) - if not isinstance(response.content, str): + response_content = message_content_to_text( + response.content, + ) + + if not response_content: return {} try: - data = cast(dict[str, object], json.loads(response.content)) + data = json.loads(response_content) except json.JSONDecodeError: - # Ignore malformed output instead of failing the pipeline. return {} - facts = cast(list[dict[str, object]], data.get("facts", [])) - for item in facts: + if not isinstance(data, dict): + return {} + + raw_facts = data.get("facts", []) + + if not isinstance(raw_facts, list): + return {} + + facts: list[Fact] = [] + + for item in raw_facts: + if not isinstance(item, dict): + continue + try: - await self.memory.remember_fact( + subject = item["subject"] + predicate = item["predicate"] + object_ = item["object"] + confidence = item.get("confidence", 1.0) + + if not isinstance(subject, str): + continue + + if not isinstance(predicate, str): + continue + + if not isinstance(object_, str): + continue + + facts.append( Fact( - subject=cast(str, item["subject"]), - predicate=cast(str, item["predicate"]), - object=cast(str, item["object"]), - confidence=float( - cast(str | float | int, item.get("confidence", 1.0)) - ), + subject=subject, + predicate=predicate, + object=object_, + confidence=float(confidence), ) ) - except (KeyError, TypeError, ValueError): - # Skip malformed facts. + + except ( + KeyError, + TypeError, + ValueError, + ): continue + if facts: + await self.memory.remember_facts(facts) + return {} diff --git a/runtime/src/orion/agent/prompts.py b/runtime/src/orion/agent/prompts.py index 3ee98a1..a48c3db 100644 --- a/runtime/src/orion/agent/prompts.py +++ b/runtime/src/orion/agent/prompts.py @@ -1,56 +1,58 @@ from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder -SYSTEM_PROMPT = """ -You are ORION, a voice assistant modeled on J.A.R.V.I.S. and F.R.I.D.A.Y. — -an unflappable, quick-witted AI aide. +# SYSTEM_PROMPT = """ +# You are ORION, a voice assistant modeled on J.A.R.V.I.S. and F.R.I.D.A.Y. — +# an unflappable, quick-witted AI aide. -Persona: -- Composed, precise, and quietly confident; a trusted right hand. -- Dry, understated wit and a touch of charm — never goofy, never verbose. -- Efficient and courteous; you get things done and report back cleanly. -- A brief wry remark is welcome, but never at the expense of clarity. +# Persona: +# - Composed, precise, and quietly confident; a trusted right hand. +# - Dry, understated wit and a touch of charm — never goofy, never verbose. +# - Efficient and courteous; you get things done and report back cleanly. +# - A brief wry remark is welcome, but never at the expense of clarity. -Getting to know the user: -- Do not assume the user's name; you serve whoever is speaking. -- If the user greets you (e.g. "hi", "hello") and you do not yet know their - name from the Retrieved Context, greet them back and politely ask for it, - for example: "Hello — may I know your good name?" -- Once the user tells you their name, remember it with the memory tools and - address them by it occasionally and naturally thereafter. -- If the Retrieved Context already contains the user's name, use it and do - not ask again. +# Getting to know the user: +# - Do not assume the user's name; you serve whoever is speaking. +# - If the user greets you (e.g. "hi", "hello") and you do not yet know their +# name from the Retrieved Context, greet them back and politely ask for it, +# for example: "Hello — may I know your good name?" +# - Once the user tells you their name, remember it with the memory tools and +# address them by it occasionally and naturally thereafter. +# - If the Retrieved Context already contains the user's name, use it and do +# not ask again. -Response Style: -- Respond naturally and conversationally, as if spoken aloud. -- Prefer a single short sentence whenever possible. -- Address the user by their name occasionally once you know it, not every line. -- Do not use Markdown, bullet points, headings, tables, or code blocks. +# Response Style: +# - Respond naturally and conversationally, as if spoken aloud. +# - Prefer a single short sentence whenever possible. +# - Address the user by their name occasionally once you know it, not every line. +# - Do not use Markdown, bullet points, headings, tables, or code blocks. -Tools and honesty (critical): -- Your file tools let you read, write, edit, move/rename, search, list, and - create folders. You CANNOT delete files. -- If asked to do something you have no tool for (for example, deleting a - file), say so plainly. NEVER claim you performed an action you did not. -- Only report a file operation as done if the tool actually returned success. -- Never invent file contents or file paths. If you need content you do not - have, read the relevant file first; if unsure where something is, list or - search before answering. -- When the result matters, read the file back to confirm before reporting. +# Tools and honesty (critical): +# - Your file tools let you read, write, edit, move/rename, search, list, and +# create folders. You CANNOT delete files. +# - If asked to do something you have no tool for (for example, deleting a +# file), say so plainly. NEVER claim you performed an action you did not. +# - Only report a file operation as done if the tool actually returned success. +# - Never invent file contents or file paths. If you need content you do not +# have, read the relevant file first; if unsure where something is, list or +# search before answering. +# - When the result matters, read the file back to confirm before reporting. -Memory: -- You have no persistent memory of your own. -- The "Retrieved Context" contains your long-term memory. -- Treat the retrieved context as the authoritative source of remembered information. -- If the retrieved context contains the answer, use it confidently. -- If the current conversation provides the answer, use it. -- If neither the retrieved context nor the current conversation contains the answer, say you do not know. -- Never invent, assume, or hallucinate facts about the user. -- Never contradict the retrieved context. +# Memory: +# - You have no persistent memory of your own. +# - The "Retrieved Context" contains your long-term memory. +# - Treat the retrieved context as the authoritative source of remembered information. +# - If the retrieved context contains the answer, use it confidently. +# - If the current conversation provides the answer, use it. +# - If neither the retrieved context nor the current conversation contains the answer, say you do not know. +# - Never invent, assume, or hallucinate facts about the user. +# - Never contradict the retrieved context. -Memory Modification: -- Use memory modification tools only when the user explicitly provides, corrects, updates, or requests deletion of stable long-term information. -- Do not store temporary conversation details unless the user indicates they should be remembered. -""" +# Memory Modification: +# - Use memory modification tools only when the user explicitly provides, corrects, updates, or requests deletion of stable long-term information. +# - Do not store temporary conversation details unless the user indicates they should be remembered. +# """ + +SYSTEM_PROMPT = """""" PROMPT = ChatPromptTemplate.from_messages( diff --git a/runtime/src/orion/events/base.py b/runtime/src/orion/events/base.py index 527dda1..9aa5663 100644 --- a/runtime/src/orion/events/base.py +++ b/runtime/src/orion/events/base.py @@ -37,10 +37,10 @@ class Event(BaseModel): event_id: UUID = Field(default_factory=uuid4) #: Correlates related events belonging to the same workflow/request. - correlation_id: UUID + correlation_id: UUID | None #: Client session that originated the event. - session_id: UUID + session_id: UUID | None #: Time at which the event was created. timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) diff --git a/runtime/src/orion/integrations/_mcp/config.py b/runtime/src/orion/integrations/_mcp/config.py index e5a397d..f709cf0 100644 --- a/runtime/src/orion/integrations/_mcp/config.py +++ b/runtime/src/orion/integrations/_mcp/config.py @@ -1,10 +1,13 @@ -# integrations/_mcp/config.py - from __future__ import annotations import json +import os from dataclasses import dataclass, field from pathlib import Path +from typing import Literal + + +MCPTransport = Literal["stdio", "http"] @dataclass(slots=True, frozen=True) @@ -14,8 +17,17 @@ class MCPServerConfig: """ name: str - command: str + + transport: MCPTransport = "stdio" + + # stdio configuration + command: str | None = None args: list[str] = field(default_factory=list) + + # HTTP configuration + url: str | None = None + + # Common configuration env: dict[str, str] = field(default_factory=dict) enabled: bool = True @@ -29,27 +41,144 @@ class MCPConfig: servers: list[MCPServerConfig] = field(default_factory=list) -def load_config(path: str | Path = "mcp.json") -> MCPConfig: +def _resolve_value( + value: str, + *, + project_root: Path, +) -> str: + """ + Resolve variables used inside MCP configuration values. + + Supported variables: + + ${PROJECT_ROOT} + $PROJECT_ROOT + + Environment variables are also expanded. """ - Build an MCPConfig from a Claude-Desktop-style JSON file: - { "mcpServers": { "": { "command": ..., "args": [...], "env": {...}, "enabled": true } } } + value = value.replace( + "${PROJECT_ROOT}", + str(project_root), + ) + + value = value.replace( + "$PROJECT_ROOT", + str(project_root), + ) + + return os.path.expandvars(value) + + +def load_config( + path: str | Path = "mcp.json", +) -> MCPConfig: + """ + Load MCP server configuration from mcp.json. + + Supported transports: + + stdio: + { + "command": "npx", + "args": [...] + } + + http: + { + "type": "http", + "url": "https://..." + } + + Configuration values may use: + + ${PROJECT_ROOT} + + which resolves to the directory containing mcp.json. """ - p = Path(path) - if not p.exists(): - return MCPConfig(servers=[]) - raw = json.loads(p.read_text()) + config_path = Path(path).resolve() + + if not config_path.exists(): + return MCPConfig() + + raw = json.loads( + config_path.read_text(), + ) + + project_root = config_path.parent + + servers: list[MCPServerConfig] = [] + + for name, spec in raw.get( + "mcpServers", + {}, + ).items(): + + transport: MCPTransport = spec.get( + "type", + "stdio", + ) + + if transport not in {"stdio", "http"}: + raise ValueError( + f"Unsupported MCP transport " + f"'{transport}' for server '{name}'." + ) + + command = spec.get("command") + + if command is not None: + command = _resolve_value( + command, + project_root=project_root, + ) + + args = [ + _resolve_value( + arg, + project_root=project_root, + ) + for arg in spec.get( + "args", + [], + ) + ] + + url = spec.get("url") + + if url is not None: + url = _resolve_value( + url, + project_root=project_root, + ) + + env = { + key: _resolve_value( + value, + project_root=project_root, + ) + for key, value in spec.get( + "env", + {}, + ).items() + } - servers = [ - MCPServerConfig( - name=name, - command=spec["command"], - args=spec.get("args", []), - env=spec.get("env", {}), - enabled=spec.get("enabled", True), + servers.append( + MCPServerConfig( + name=name, + transport=transport, + command=command, + args=args, + url=url, + env=env, + enabled=spec.get( + "enabled", + True, + ), + ) ) - for name, spec in raw.get("mcpServers", {}).items() - ] - return MCPConfig(servers=servers) + return MCPConfig( + servers=servers, + ) diff --git a/runtime/src/orion/integrations/_mcp/events.py b/runtime/src/orion/integrations/_mcp/events.py new file mode 100644 index 0000000..995de18 --- /dev/null +++ b/runtime/src/orion/integrations/_mcp/events.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from orion.events.base import Event, EventStatus + + +# ========================================================== +# Server Lifecycle +# ========================================================== + + +class MCPServerStartedEvent(Event): + """ + Published when an MCP server successfully connects. + """ + + status: EventStatus = EventStatus.SUCCESS + + server_name: str + transport: str + + +class MCPServerStartupFailedEvent(Event): + """ + Published when an MCP server fails to start. + """ + + status: EventStatus = EventStatus.ERROR + + server_name: str + transport: str + error: str + + +class MCPServerStoppedEvent(Event): + """ + Published when an MCP server is shut down successfully. + """ + + status: EventStatus = EventStatus.SUCCESS + + server_name: str + + +class MCPServerShutdownFailedEvent(Event): + """ + Published when an MCP server fails to shut down. + """ + + status: EventStatus = EventStatus.ERROR + + server_name: str + error: str + + +# ========================================================== +# Tool Discovery +# ========================================================== + + +class MCPToolsDiscoveredEvent(Event): + """ + Published after an MCP server's tools have been + successfully discovered. + """ + + status: EventStatus = EventStatus.SUCCESS + + server_name: str + tool_names: list[str] + tool_count: int + + +class MCPToolsDiscoveryFailedEvent(Event): + """ + Published when tool discovery fails for an MCP server. + """ + + status: EventStatus = EventStatus.ERROR + + server_name: str + error: str + + +# ========================================================== +# Tool Execution +# ========================================================== + + +class MCPToolCalledEvent(Event): + """ + Published immediately before an MCP tool is executed. + + This is request-scoped and therefore carries the + session and correlation identifiers inherited from Event. + """ + + status: EventStatus = EventStatus.INFO + + server_name: str + tool_name: str + + +class MCPToolCompletedEvent(Event): + """ + Published after an MCP tool completes successfully. + + This is request-scoped. + """ + + status: EventStatus = EventStatus.SUCCESS + + server_name: str + tool_name: str + + +class MCPToolFailedEvent(Event): + """ + Published when an MCP tool execution fails. + + This is request-scoped. + """ + + status: EventStatus = EventStatus.ERROR + + server_name: str + tool_name: str + error: str diff --git a/runtime/src/orion/integrations/_mcp/langchain.py b/runtime/src/orion/integrations/_mcp/langchain.py new file mode 100644 index 0000000..b61297f --- /dev/null +++ b/runtime/src/orion/integrations/_mcp/langchain.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from typing import Any +from uuid import UUID + +from langchain_core.tools import BaseTool +from pydantic import ConfigDict, PrivateAttr + +from orion.integrations._mcp.manager import MCPManager + + +class MCPTool(BaseTool): + """ + Thin LangChain adapter around an MCP tool. + + MCPManager owns: + - MCP server lifecycle + - tool discovery + - MCP protocol communication + - tool execution + + MCPTool only adapts an MCP tool to the LangChain/LangGraph + tool interface. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + _manager: MCPManager = PrivateAttr() + _parameters: dict[str, Any] = PrivateAttr() + _session_id: UUID = PrivateAttr() + _correlation_id: UUID = PrivateAttr() + + def __init__( + self, + manager: MCPManager, + *, + name: str, + description: str, + parameters: dict[str, Any], + session_id: UUID, + correlation_id: UUID, + ) -> None: + super().__init__( + name=name, + description=description, + ) + + self._manager = manager + self._parameters = parameters + self._session_id = session_id + self._correlation_id = correlation_id + + @property + def args(self) -> dict[str, Any]: + """ + Return the original MCP JSON schema unchanged. + + LangGraph/LangChain can use this to expose the exact + schema supplied by the MCP server. + """ + return self._parameters + + def _run( + self, + **kwargs: Any, + ) -> str: + raise NotImplementedError( + "MCPTool only supports asynchronous execution." + ) + + async def _arun( + self, + **kwargs: Any, + ) -> str: + return await self._manager.call_tool( + self.name, + kwargs, + session_id=self._session_id, + correlation_id=self._correlation_id, + ) + + +def create_mcp_tools( + manager: MCPManager, + *, + session_id: UUID, + correlation_id: UUID, +) -> list[MCPTool]: + """ + Convert MCPManager's discovered tools into LangChain tools. + + The MCP JSON schema is preserved exactly. No Pydantic model + or schema transformation is performed. + """ + + tools: list[MCPTool] = [] + + for schema in manager.tools: + function = schema["function"] + + name = function["name"] + description = function.get("description", "") + + parameters = function.get( + "parameters", + { + "properties": {}, + "type": "object", + }, + ) + + tools.append( + MCPTool( + manager, + name=name, + description=description, + parameters=parameters, + session_id=session_id, + correlation_id=correlation_id, + ) + ) + + return tools diff --git a/runtime/src/orion/integrations/_mcp/langchain_tools.py b/runtime/src/orion/integrations/_mcp/langchain_tools.py deleted file mode 100644 index 2ce455d..0000000 --- a/runtime/src/orion/integrations/_mcp/langchain_tools.py +++ /dev/null @@ -1,113 +0,0 @@ -# integrations/_mcp/langchain_tools.py - -from __future__ import annotations - -import json -import os -from pathlib import Path - -from langchain_core.tools import BaseTool -from langchain_mcp_adapters.client import MultiServerMCPClient -from langchain_mcp_adapters.sessions import Connection, StdioConnection - - -def _project_root() -> Path: - """ - Determine the ORION project root deterministically, independent of the - current working directory: walk up from this file until a project marker - (pyproject.toml or .git) is found. Falls back to the cwd. - """ - here = Path(__file__).resolve() - for parent in (here, *here.parents): - if (parent / "pyproject.toml").exists() or (parent / ".git").exists(): - return parent - return Path.cwd() - - -def _resolve(value: str) -> str: - """ - Expand config placeholders so paths are never hardcoded to one machine: - - - ``${PROJECT_ROOT}`` -> the ORION project root, found from the code's own - location (correct no matter which directory you launch from). - - ``${CWD}`` / ``${PWD}`` -> the current working directory. - - ``$VAR`` / ``${VAR}`` -> the matching environment variable. - - This lets `mcp.json` ship a portable value that resolves to whatever - machine and directory the product actually runs in. - """ - root = str(_project_root()) - cwd = str(Path.cwd()) - value = ( - value.replace("${PROJECT_ROOT}", root) - .replace("${CWD}", cwd) - .replace("${PWD}", cwd) - ) - return os.path.expandvars(value) - - -async def load_mcp_tools(config_path: str | Path = "mcp.json") -> list[BaseTool]: - """ - Load tools from every enabled MCP server declared in `mcp.json` and return - them as LangChain BaseTools, ready for `llm.bind_tools(...)` and `ToolNode`. - - Config format (Claude-Desktop style): - - { - "mcpServers": { - "": { - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-filesystem", "${CWD}"], - "env": {"KEY": "value"}, - "enabled": true, - "tools": ["read_text_file", "list_directory"] - } - } - } - - Placeholders in "args" (``${CWD}`` / ``${PWD}`` / ``$VAR``) are expanded at - runtime, so paths adapt to the machine ORION runs on instead of being - hardcoded. - - The optional per-server "tools" allowlist keeps only the named tools. - Every tool schema is sent to the LLM on *every* call, so exposing only - the few you actually need meaningfully cuts prompt tokens. When no - server declares "tools", all discovered tools are used. - - Returns an empty list if the config file is missing or has no enabled - servers, so the agent degrades gracefully to its built-in OrionTools. - """ - path = Path(config_path) - if not path.exists(): - return [] - - raw = json.loads(path.read_text()) - - connections: dict[str, Connection] = {} - allowlist: set[str] = set() - - for name, spec in raw.get("mcpServers", {}).items(): - if not spec.get("enabled", True): - continue - - connection: Connection = StdioConnection( - transport="stdio", - command=spec["command"], - args=[_resolve(str(arg)) for arg in spec.get("args", [])], - env=spec.get("env") or None, - ) - - connections[name] = connection - allowlist.update(spec.get("tools", []) or []) - - if not connections: - return [] - - client = MultiServerMCPClient(connections) - tools = await client.get_tools() - - # If any server declared an allowlist, keep only those tools. - if allowlist: - tools = [tool for tool in tools if tool.name in allowlist] - - return tools diff --git a/runtime/src/orion/integrations/_mcp/manager.py b/runtime/src/orion/integrations/_mcp/manager.py index 4761522..51158b8 100644 --- a/runtime/src/orion/integrations/_mcp/manager.py +++ b/runtime/src/orion/integrations/_mcp/manager.py @@ -1,80 +1,478 @@ -# integrations/_mcp/manager.py +# runtime/src/orion/integrations/_mcp/manager.py from __future__ import annotations -from typing import cast +from typing import Any, cast +from uuid import UUID +from rich.console import Console +from typing_extensions import override + +from orion.bus.event_bus import EventBus from orion.integrations._mcp.config import MCPConfig -from orion.integrations._mcp.server import MCPServer from orion.integrations._mcp.discovery import mcp_tools_to_openai +from orion.integrations._mcp.events import ( + MCPServerShutdownFailedEvent, + MCPServerStartedEvent, + MCPServerStartupFailedEvent, + MCPServerStoppedEvent, + MCPToolCalledEvent, + MCPToolCompletedEvent, + MCPToolFailedEvent, + MCPToolsDiscoveredEvent, + MCPToolsDiscoveryFailedEvent, +) +from orion.integrations._mcp.server import MCPServer +from orion.runtime.lifecycle import Lifecycle + +console = Console() -class MCPManager: + +class MCPManager(Lifecycle): """ - Manages the lifecycle of all configured MCP servers and exposes their - tools to the agent in an OpenAI/Groq-compatible format. + Manages all configured MCP servers. + + Responsibilities: + - Start and stop MCP servers. + - Discover tools exposed by connected servers. + - Maintain tool -> server routing. + - Execute MCP tools. + - Expose discovered tool schemas to the agent. + - Publish MCP lifecycle and execution events. + + MCPManager is transport- and server-agnostic. """ + SOURCE = "mcp_manager" + def __init__( self, config: MCPConfig, ) -> None: self._config = config + + # EventBus is a process-wide singleton. + # It must already be initialized by the application runtime. + self._bus = EventBus() + self._servers: dict[str, MCPServer] = {} + + # MCP tool name -> owning MCP server name. self._tool_routing: dict[str, str] = {} - self._tools: list[dict[str, object]] = [] + # OpenAI/Groq-compatible tool schemas. + self._tools: list[dict[str, Any]] = [] + + self._started = False + + # ========================================================== + # Properties + # ========================================================== + + @property + def servers(self) -> dict[str, MCPServer]: + """ + Return all successfully connected MCP servers. + """ + return self._servers + + @property + def tools(self) -> list[dict[str, Any]]: + """ + Return all discovered MCP tools in + OpenAI/Groq-compatible format. + """ + return self._tools + + @property + def started(self) -> bool: + """ + Whether the manager has been started. + """ + return self._started + + # ========================================================== + # Lifecycle + # ========================================================== + + @override async def startup(self) -> None: """ - Start every enabled MCP server and collect its tools. - A failing server is skipped, not fatal to the pipeline. + Start all enabled MCP servers and discover their tools. + + Failure of one MCP server does not prevent other configured + servers from starting. """ + + if self._started: + return + + console.print( + "[cyan]Starting MCP servers...[/]" + ) + for server_config in self._config.servers: + if not server_config.enabled: + console.print( + f"[dim]Skipping disabled MCP server " + f"'{server_config.name}'.[/]" + ) continue server = MCPServer(server_config) + + # -------------------------------------------------- + # Start server + # -------------------------------------------------- + try: await server.startup() - except Exception as e: - print(f"[mcp] failed to start '{server_config.name}': {e}") + + except Exception as exc: + console.print( + f"[red]✗ MCP server " + f"'{server_config.name}' failed to start:[/] " + f"{exc}" + ) + + await self._bus.publish( + MCPServerStartupFailedEvent( + correlation_id=None, + session_id=None, + source=self.SOURCE, + message="MCP server startup failed.", + server_name=server_config.name, + transport=server_config.transport, + error=str(exc), + ) + ) + continue self._servers[server.name] = server - result = await server.list_tools() - for schema in mcp_tools_to_openai(result): - tool_name = schema["function"]["name"] + console.print( + f"[green]✓ MCP server " + f"'{server.name}' connected " + f"({server_config.transport}).[/]" + ) + + await self._bus.publish( + MCPServerStartedEvent( + correlation_id=None, + session_id=None, + source=self.SOURCE, + message="MCP server started.", + server_name=server.name, + transport=server_config.transport, + ) + ) + + # -------------------------------------------------- + # Discover tools + # -------------------------------------------------- + + try: + result = await server.list_tools() + + schemas = mcp_tools_to_openai(result) + + except Exception as exc: + console.print( + f"[red]✗ Failed to discover tools " + f"from '{server.name}':[/] {exc}" + ) + + await self._bus.publish( + MCPToolsDiscoveryFailedEvent( + correlation_id=None, + session_id=None, + source=self.SOURCE, + message="MCP tool discovery failed.", + server_name=server.name, + error=str(exc), + ) + ) + + continue + + discovered_names: list[str] = [] + + for schema in schemas: + + function = cast( + dict[str, Any], + schema.get("function", {}), + ) + + tool_name = function.get("name") + + if not isinstance(tool_name, str): + continue + + # -------------------------------------------------- + # Duplicate protection + # -------------------------------------------------- + + if tool_name in self._tool_routing: + existing_server = self._tool_routing[tool_name] + + console.print( + f"[yellow]⚠ MCP tool " + f"'{tool_name}' from '{server.name}' " + f"was ignored because it is already owned " + f"by '{existing_server}'.[/]" + ) + + continue + self._tool_routing[tool_name] = server.name self._tools.append(schema) + discovered_names.append(tool_name) - @property - def tools(self) -> list[dict[str, object]]: - """OpenAI/Groq tool schemas for every connected server.""" - return self._tools + console.print( + f"[green]✓ Discovered " + f"{len(discovered_names)} tools " + f"from '{server.name}'.[/]" + ) + + await self._bus.publish( + MCPToolsDiscoveredEvent( + correlation_id=None, + session_id=None, + source=self.SOURCE, + message="MCP tools discovered.", + server_name=server.name, + tool_names=discovered_names, + tool_count=len(discovered_names), + ) + ) + + self._started = True + + console.print( + f"[green]MCP startup complete: " + f"{len(self._servers)} server(s), " + f"{len(self._tools)} tool(s).[/]" + ) + + @override + async def shutdown(self) -> None: + """ + Gracefully shut down all connected MCP servers. + """ + + if not self._started and not self._servers: + return + + console.print( + "[cyan]Stopping MCP servers...[/]" + ) + + for server_name, server in list( + self._servers.items() + ): + try: + await server.shutdown() + + except Exception as exc: + console.print( + f"[red]✗ MCP server " + f"'{server_name}' failed to shut down:[/] " + f"{exc}" + ) + + await self._bus.publish( + MCPServerShutdownFailedEvent( + correlation_id=None, + session_id=None, + source=self.SOURCE, + message="MCP server shutdown failed.", + server_name=server_name, + error=str(exc), + ) + ) + + continue + + console.print( + f"[green]✓ MCP server " + f"'{server_name}' stopped.[/]" + ) + + await self._bus.publish( + MCPServerStoppedEvent( + correlation_id=None, + session_id=None, + source=self.SOURCE, + message="MCP server stopped.", + server_name=server_name, + ) + ) + + self._servers.clear() + self._tool_routing.clear() + self._tools.clear() + + self._started = False + + console.print( + "[green]MCP shutdown complete.[/]" + ) + + # ========================================================== + # Server Access + # ========================================================== + + def server( + self, + name: str, + ) -> MCPServer: + """ + Return a connected MCP server by name. + """ + + try: + return self._servers[name] + + except KeyError: + raise ValueError( + f"Unknown MCP server: {name}" + ) from None + + # ========================================================== + # Tool Execution + # ========================================================== + + async def call_tool_raw( + self, + name: str, + arguments: dict[str, object], + *, + session_id: UUID, + correlation_id: UUID, + ) -> object: + """ + Route an MCP tool call to its owning server. + + Returns the raw MCP result. + """ - async def call_tool(self, name: str, arguments: dict[str, object]) -> str: - """Route a tool call to the server that owns it and flatten the result to text.""" server_name = self._tool_routing.get(name) + if server_name is None: - return f"Error: unknown tool '{name}'" + raise ValueError( + f"Unknown MCP tool: {name}" + ) - result = await self._servers[server_name].call_tool(name, arguments) + server = self._servers.get(server_name) - content = cast(list[object], getattr(result, "content", [])) - parts: list[str] = [] - for item in content: - text = cast(str | None, getattr(item, "text", None)) - if text is not None: - parts.append(text) + if server is None: + raise RuntimeError( + f"MCP server '{server_name}' " + f"owning tool '{name}' is not connected" + ) - return "\n".join(parts) if parts else "(no output)" + await self._bus.publish( + MCPToolCalledEvent( + correlation_id=correlation_id, + session_id=session_id, + source=self.SOURCE, + message="MCP tool execution started.", + server_name=server_name, + tool_name=name, + ) + ) - async def shutdown(self) -> None: + console.print( + f"[cyan]→ MCP tool " + f"'{name}' " + f"({server_name})[/]" + ) + + try: + result = await server.call_tool( + name, + arguments, + ) + + except Exception as exc: + console.print( + f"[red]✗ MCP tool " + f"'{name}' failed:[/] {exc}" + ) + + await self._bus.publish( + MCPToolFailedEvent( + correlation_id=correlation_id, + session_id=session_id, + source=self.SOURCE, + message="MCP tool execution failed.", + server_name=server_name, + tool_name=name, + error=str(exc), + ) + ) + + raise + + await self._bus.publish( + MCPToolCompletedEvent( + correlation_id=correlation_id, + session_id=session_id, + source=self.SOURCE, + message="MCP tool execution completed.", + server_name=server_name, + tool_name=name, + ) + ) + + console.print( + f"[green]✓ MCP tool " + f"'{name}' completed.[/]" + ) + + return result + + async def call_tool( + self, + name: str, + arguments: dict[str, object], + *, + session_id: UUID, + correlation_id: UUID, + ) -> str: """ - Shuts down all configured MCP servers. + Execute an MCP tool and flatten textual content into a string. """ - for server in self._servers.values(): - await server.shutdown() - self._servers.clear() + result = await self.call_tool_raw( + name, + arguments, + session_id=session_id, + correlation_id=correlation_id, + ) + + content = cast( + list[object], + getattr(result, "content", []), + ) + + parts: list[str] = [] + + for item in content: + text = cast( + str | None, + getattr(item, "text", None), + ) + + if text is not None: + parts.append(text) + + return ( + "\n".join(parts) + if parts + else "(no output)" + ) diff --git a/runtime/src/orion/integrations/_mcp/server.py b/runtime/src/orion/integrations/_mcp/server.py index 892b6ec..1d917d9 100644 --- a/runtime/src/orion/integrations/_mcp/server.py +++ b/runtime/src/orion/integrations/_mcp/server.py @@ -3,6 +3,7 @@ from __future__ import annotations from contextlib import AsyncExitStack +from typing import Any from mcp import ClientSession @@ -11,11 +12,14 @@ stdio_client, get_default_environment, ) +from typing_extensions import override +from mcp.client.streamable_http import streamable_http_client from orion.integrations._mcp.config import MCPServerConfig +from orion.runtime.lifecycle import Lifecycle -class MCPServer: +class MCPServer(Lifecycle): """ Represents a single running MCP server. """ @@ -30,6 +34,11 @@ def __init__( self._session: ClientSession | None = None self._connected: bool = False + # ================================================= + # Properties + # ================================================= + + @property def name(self) -> str: return self.config.name @@ -38,21 +47,128 @@ def name(self) -> str: def connected(self) -> bool: return self._connected + + # ================================================= + # Lifecycle + # ================================================= + + + @override async def startup(self) -> None: + if self._connected: + return + + if self.config.transport == "stdio": + await self._startup_stdio() + elif self.config.transport == "http": + await self._startup_http() + + else: + raise ValueError( + f"Unsupported MCP transport " + f"'{self.config.transport}'" + f"for server '{self.name}'" + ) + + self._connected = True + + @override + async def shutdown(self) -> None: """ - Launch the server and establish an MCP session. + Gracefully shut down the MCP server. + + Internal state is reset even if the underlying transport + fails to close cleanly. """ + if not self._connected: + return + + try: + await self._stack.aclose() + + finally: + self._session = None + self._connected = False + self._stack = AsyncExitStack() + + # ================================================= + # Startup stdio + # ================================================= + + async def _startup_stdio(self) -> None: + """ + Start the stdio stream for the MCP server. + """ + + if not self.config.command: + raise ValueError( + f"MCP Server '{self.name}' uses stdio " + "but no command was configured" + ) + params = StdioServerParameters( command=self.config.command, args=self.config.args, - env={**get_default_environment(), **self.config.env}, + env= { + **get_default_environment(), + **(self.config.env or {}), + } ) - read_stream, write_stream = await self._stack.enter_async_context( - stdio_client(params) + read_stream, write_stream = ( + await self._stack.enter_async_context( + stdio_client(params) + ) + ) + + self._session = await self._create_session( + read_stream, + write_stream, + ) + + # ================================================= + # Startup HTTP + # ================================================= + + + async def _startup_http(self) -> None: + """ + Start an MCP connection over HTTP. + + The concrete HTTP transport implementation will be added + here using the MCP SDK's HTTP client transport. + """ + + if not self.config.url: + raise ValueError( + f"MCP Server '{self.name}' uses HTTP " + "but no URL was configured" + ) + + read_stream, write_stream, _ = ( + await self._stack.enter_async_context( + streamable_http_client( + self.config.url, + ) + ) ) + self._session = await self._create_session(read_stream, write_stream) + + # ================================================= + # Session Management + # ================================================= + + async def _create_session( + self, + read_stream, + write_stream, + ) -> ClientSession: + """ + Create and initialize a new MCP client session. + """ + session = await self._stack.enter_async_context( ClientSession( read_stream, @@ -60,36 +176,34 @@ async def startup(self) -> None: ) ) - _ = await session.initialize() + await session.initialize() + return session - self._session = session - self._connected = True + # ================================================= + # MCP Operations + # ================================================= - async def shutdown(self) -> None: + async def list_tools(self): """ - Close the MCP session. + List all available tools on the MCP server. """ - await self._stack.aclose() - - self._session = None - self._connected = False - - async def list_tools(self): if self._session is None: - raise RuntimeError("MCP server is not connected.") + raise RuntimeError( + f"MCP Server '{self.name}' is not connected" + ) return await self._session.list_tools() - async def call_tool( - self, - name: str, - arguments: dict[str, object], - ) -> object: + + async def call_tool(self, name: str, arguments: dict[str, object]): + """ + Call a tool on the MCP server. + """ + if self._session is None: - raise RuntimeError("MCP server is not connected.") + raise RuntimeError( + f"MCP Server '{self.name}' is not connected" + ) - return await self._session.call_tool( - name, - arguments, - ) + return await self._session.call_tool(name, arguments) diff --git a/runtime/src/orion/llm/__init__.py b/runtime/src/orion/llm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/runtime/src/orion/llm/config.py b/runtime/src/orion/llm/config.py new file mode 100644 index 0000000..f31c42f --- /dev/null +++ b/runtime/src/orion/llm/config.py @@ -0,0 +1,57 @@ +from dataclasses import dataclass, field +import os + +from dotenv import load_dotenv + + +load_dotenv() + + +@dataclass(slots=True) +class LLMConfig: + """ + Configuration for the ORION language model. + + The provider determines which LangChain chat model is + constructed by the LLM factory. + """ + + # ========================================================== + # Provider + # ========================================================== + + provider: str = field( + default_factory=lambda: os.getenv( + "ORION_LLM_PROVIDER", + "gemini", + ) + ) + + # ========================================================== + # Model + # ========================================================== + + model: str = field( + default_factory=lambda: os.getenv( + "ORION_LLM_MODEL", + "", + ) + ) + + # ========================================================== + # API Keys + # ========================================================== + + gemini_api_key: str = field( + default_factory=lambda: os.getenv( + "GEMINI_API_KEY", + "", + ) + ) + + groq_api_key: str = field( + default_factory=lambda: os.getenv( + "GROQ_API_KEY", + "", + ) + ) diff --git a/runtime/src/orion/llm/factory.py b/runtime/src/orion/llm/factory.py new file mode 100644 index 0000000..5190d9c --- /dev/null +++ b/runtime/src/orion/llm/factory.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from langchain_core.language_models import BaseChatModel + +from orion.llm.config import LLMConfig +from orion.llm.interfaces.provider import LLMProvider +from orion.llm.providers.gemini import GeminiProvider +from orion.llm.providers.groq import GroqProvider + + +@dataclass(slots=True) +class LLMProviders: + """ + Container for LLM providers. + """ + + provider: LLMProvider + + +class LLMFactory: + """ + Creates the configured LLM provider. + """ + + @staticmethod + def create( + config: LLMConfig, + ) -> LLMProviders: + + if config.provider == "gemini": + provider = GeminiProvider(config) + + elif config.provider == "groq": + provider = GroqProvider(config) + + else: + raise ValueError( + f"Unsupported LLM provider: {config.provider}" + ) + + return LLMProviders( + provider=provider, + ) diff --git a/runtime/src/orion/llm/interfaces/provider.py b/runtime/src/orion/llm/interfaces/provider.py new file mode 100644 index 0000000..2be5ab7 --- /dev/null +++ b/runtime/src/orion/llm/interfaces/provider.py @@ -0,0 +1,16 @@ +from abc import ABC, abstractmethod + +from langchain_core.language_models import BaseChatModel + + +class LLMProvider(ABC): + """ + Abstract provider for creating LangChain chat models. + """ + + @abstractmethod + def create(self) -> BaseChatModel: + """ + Create and return the configured chat model. + """ + raise NotImplementedError diff --git a/runtime/src/orion/llm/providers/gemini.py b/runtime/src/orion/llm/providers/gemini.py new file mode 100644 index 0000000..8cffb28 --- /dev/null +++ b/runtime/src/orion/llm/providers/gemini.py @@ -0,0 +1,35 @@ +from langchain_core.language_models import BaseChatModel +from langchain_google_genai import ChatGoogleGenerativeAI +from typing_extensions import override + +from orion.llm.config import LLMConfig +from orion.llm.interfaces.provider import LLMProvider + + +class GeminiProvider(LLMProvider): + """ + Google Gemini LLM provider. + """ + + def __init__( + self, + config: LLMConfig, + ) -> None: + self.config = config + + @override + def create(self) -> BaseChatModel: + if not self.config.gemini_api_key: + raise ValueError( + "Gemini API key not set" + ) + + if not self.config.model: + raise ValueError( + "Gemini model not set" + ) + + return ChatGoogleGenerativeAI( + model=self.config.model, + google_api_key=self.config.gemini_api_key, + ) diff --git a/runtime/src/orion/llm/providers/groq.py b/runtime/src/orion/llm/providers/groq.py new file mode 100644 index 0000000..7c0f6dd --- /dev/null +++ b/runtime/src/orion/llm/providers/groq.py @@ -0,0 +1,37 @@ +from langchain_core.language_models import BaseChatModel +from langchain_groq import ChatGroq +from pydantic import SecretStr +from typing_extensions import override +from typing import cast + +from orion.llm.config import LLMConfig +from orion.llm.interfaces.provider import LLMProvider + + +class GroqProvider(LLMProvider): + """ + Groq LLM provider. + """ + + def __init__( + self, + config: LLMConfig, + ) -> None: + self.config = config + + @override + def create(self) -> BaseChatModel: + if not self.config.groq_api_key: + raise ValueError( + "Groq API key not set" + ) + + if not self.config.model: + raise ValueError( + "Groq model not set" + ) + + return ChatGroq( + model=self.config.model, + api_key=cast(SecretStr, self.config.groq_api_key), + ) diff --git a/runtime/src/orion/llm/utils.py b/runtime/src/orion/llm/utils.py new file mode 100644 index 0000000..b4fc99d --- /dev/null +++ b/runtime/src/orion/llm/utils.py @@ -0,0 +1,33 @@ +from __future__ import annotations + + +def message_content_to_text(content: object) -> str: + """ + Convert LangChain message content into plain text. + + Supports both plain string content and + structured content blocks. + """ + + if isinstance(content, str): + return content + + if not isinstance(content, list): + return "" + + parts: list[str] = [] + + for block in content: + if isinstance(block, str): + parts.append(block) + continue + + if not isinstance(block, dict): + continue + + text = block.get("text") + + if isinstance(text, str): + parts.append(text) + + return "".join(parts) diff --git a/runtime/src/orion/memory/config.py b/runtime/src/orion/memory/config.py index 52f4b9a..f5ab345 100644 --- a/runtime/src/orion/memory/config.py +++ b/runtime/src/orion/memory/config.py @@ -1,15 +1,84 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field +import os + +from dotenv import load_dotenv + + +load_dotenv() @dataclass(slots=True) class MemoryConfig: - sqlite_db: str = "orion.db" + """ + Configuration for the ORION memory subsystem. + + Values are loaded from environment variables when available, + with sensible local-development defaults. + """ + + # ========================================================== + # SQLite + # ========================================================== + + sqlite_db: str = field( + default_factory=lambda: os.getenv( + "ORION_SQLITE_DB", + "orion.db", + ) + ) + + # ========================================================== + # Embeddings + # ========================================================== + + embedding_model: str = field( + default_factory=lambda: os.getenv( + "ORION_EMBEDDING_MODEL", + "BAAI/bge-small-en-v1.5", + ) + ) + + # ========================================================== + # Qdrant + # ========================================================== + + qdrant_host: str = field( + default_factory=lambda: os.getenv( + "QDRANT_HOST", + "localhost", + ) + ) + + qdrant_port: int = field( + default_factory=lambda: int( + os.getenv( + "QDRANT_PORT", + "6333", + ) + ) + ) + + # ========================================================== + # Neo4j + # ========================================================== - embedding_model: str = "BAAI/bge-small-en-v1.5" + neo4j_uri: str = field( + default_factory=lambda: os.getenv( + "NEO4J_URI", + "bolt://localhost:7687", + ) + ) - qdrant_host: str = "localhost" - qdrant_port: int = 6333 + neo4j_username: str = field( + default_factory=lambda: os.getenv( + "NEO4J_USERNAME", + "neo4j", + ) + ) - neo4j_uri: str = "bolt://localhost:7687" - neo4j_username: str = "neo4j" - neo4j_password: str = "orion123" + neo4j_password: str = field( + default_factory=lambda: os.getenv( + "NEO4J_PASSWORD", + "", + ) + ) diff --git a/runtime/src/orion/memory/interfaces/graph.py b/runtime/src/orion/memory/interfaces/graph.py index f67ece0..8068dd5 100644 --- a/runtime/src/orion/memory/interfaces/graph.py +++ b/runtime/src/orion/memory/interfaces/graph.py @@ -1,65 +1,91 @@ from abc import ABC, abstractmethod -from orion.memory.models import Entity, Fact, Relationship +from orion.memory.models import Entity, Fact, GraphSchema class KnowledgeGraph(ABC): """ - Abstract Knowledge Graph. + Abstract knowledge graph. - Graph is reponsible only for storing and - retrieving structured facts + Responsible only for storing and retrieving structured knowledge. - It never dedicates what should be remastered - The LLM controls all graph mutations through tools + Implementations may use Neo4j, an MCP server, an in-memory graph, + or any other graph backend. """ + # ========================================================== + # Lifecycle + # ========================================================== + @abstractmethod async def startup(self) -> None: """ - Initialize the graph backend + Initialize the graph backend. """ raise NotImplementedError @abstractmethod async def shutdown(self) -> None: """ - Gracefully shutdown the graph backend + Gracefully shutdown the graph backend. """ raise NotImplementedError + # ========================================================== + # Storage + # ========================================================== + @abstractmethod - async def add_entity(self, entity: Entity) -> None: + async def add_fact( + self, + fact: Fact, + ) -> None: """ - Insert or Update an Entity + Store a single fact. """ raise NotImplementedError @abstractmethod - async def add_relationship(self, relationship: Relationship) -> None: + async def add_facts( + self, + facts: list[Fact], + ) -> None: """ - Insert or update a Relationship + Store multiple facts efficiently. """ raise NotImplementedError @abstractmethod - async def add_fact(self, fact: Fact) -> None: + async def remove_fact( + self, + fact: Fact, + ) -> None: """ - Convinience API for adding a fast triple + Remove a single fact. """ raise NotImplementedError @abstractmethod - async def remove_fact(self, fact: Fact) -> None: + async def remove_facts( + self, + facts: list[Fact], + ) -> None: """ - Remove a fact from the graph + Remove multiple facts. """ raise NotImplementedError + # ========================================================== + # Retrieval + # ========================================================== + @abstractmethod - async def query(self, query: str) -> list[Fact]: + async def search_facts( + self, + query: str, + ) -> list[Fact]: """ - Query the knowledge Graph + Search for facts relevant to a query or entity. """ raise NotImplementedError @@ -71,20 +97,31 @@ async def related_entities( depth: int = 1, ) -> list[Entity]: """ - Retrieve entities related to the given entity. + Retrieve entities connected to the supplied entity. """ raise NotImplementedError + @abstractmethod + async def get_schema(self) -> GraphSchema: + """ + Retrieve the graph schema. + """ + raise NotImplementedError + + # ========================================================== + # Maintenance + # ========================================================== + @abstractmethod async def clear(self) -> None: """ - Remove all graph data. + Remove every stored entity and fact. """ raise NotImplementedError @abstractmethod async def count(self) -> int: """ - Return the number of stored facts. + Return the total number of stored facts. """ raise NotImplementedError diff --git a/runtime/src/orion/memory/models.py b/runtime/src/orion/memory/models.py index 02741e6..848810f 100644 --- a/runtime/src/orion/memory/models.py +++ b/runtime/src/orion/memory/models.py @@ -136,6 +136,18 @@ class MemoryContext: tags: list[str] = field(default_factory=list) +@dataclass(slots=True) +class GraphSchema: + """ + Schema describing the structure of the knowledge graph + + This is used to guide llms to towards reusing existing labels and + relationships types instead of inventing new ones. + """ + + labels: list[str] = field(default_factory=list) + + relationship_types: list[str] = field(default_factory=list) @dataclass(slots=True) class RetrievedContext: @@ -146,7 +158,9 @@ class RetrievedContext: summary: SummaryMemory | None = None # semantic memories episodes: list[ConversationEpisode] = field(default_factory=list) - # graph + # graph schema + graph_schema: GraphSchema | None = None + # graph facts facts: list[Fact] = field(default_factory=list) # recent chat history recent_messages: list[ConversationEpisode] = field(default_factory=list) diff --git a/runtime/src/orion/memory/providers/graph/neo4j.py b/runtime/src/orion/memory/providers/graph/neo4j.py index 441de6c..ea3dc3b 100644 --- a/runtime/src/orion/memory/providers/graph/neo4j.py +++ b/runtime/src/orion/memory/providers/graph/neo4j.py @@ -1,120 +1,163 @@ +from __future__ import annotations + +import os +import re from typing import LiteralString, cast +from dotenv import load_dotenv from neo4j import AsyncDriver, AsyncGraphDatabase +from typing_extensions import override from orion.memory.interfaces.graph import KnowledgeGraph -from orion.memory.models import Entity, Fact, Relationship +from orion.memory.models import Entity, Fact, GraphSchema + +load_dotenv() + + +def normalize_predicate(predicate: str) -> str: + """ + Normalize an LLM-generated fact predicate into a valid + Neo4j relationship type. + + Examples: + "uses operating system" + -> "USES_OPERATING_SYSTEM" + + "works at" + -> "WORKS_AT" + + "favorite language" + -> "FAVORITE_LANGUAGE" + """ + + normalized = predicate.strip().upper() + + # Replace every non-alphanumeric character with "_". + normalized = re.sub( + r"[^A-Z0-9]+", + "_", + normalized, + ) + + # Collapse repeated underscores. + normalized = re.sub( + r"_+", + "_", + normalized, + ) + + # Remove leading/trailing underscores. + normalized = normalized.strip("_") + + if not normalized: + raise ValueError( + f"Invalid empty predicate: {predicate!r}" + ) + + return normalized class Neo4jKnowledgeGraph(KnowledgeGraph): """ - Neo4j backed implementation of the ORION knowledge graph. + Neo4j-backed implementation of the ORION knowledge graph. - This class is purely responsible for persistence. - The LLM decides what gets inserted or removed. + Responsible only for persistence and retrieval of structured + knowledge. The caller decides what knowledge should be stored + or retrieved. """ def __init__( self, - uri: str = "bolt://localhost:7687", - username: str = "neo4j", - password: str = "orion123", + uri: str, + username: str, + password: str, ) -> None: self.uri = uri self.username = username self.password = password + self.driver: AsyncDriver | None = None + # ========================================================== + # Lifecycle + # ========================================================== + + @override async def startup(self) -> None: self.driver = AsyncGraphDatabase.driver( self.uri, - auth=(self.username, self.password), + auth=( + self.username, + self.password, + ), ) + await self.driver.verify_connectivity() + @override async def shutdown(self) -> None: if self.driver is not None: await self.driver.close() + self.driver = None - async def add_entity( - self, - entity: Entity, - ) -> None: - assert self.driver is not None - - query = """ - MERGE (e:Entity {name: $name}) - SET e.label = $label - """ - - async with self.driver.session() as session: - await session.run( - query, - name=entity.name, - label=entity.label, - ) - - async def add_relationship( - self, - relationship: Relationship, - ) -> None: - assert self.driver is not None - - query = cast( - LiteralString, - f""" - MATCH (a:Entity {{name: $source}}) - MATCH (b:Entity {{name: $target}}) - MERGE (a)-[r:{relationship.predicate}]->(b) - SET r.confidence = $confidence - """, - ) - - async with self.driver.session() as session: - await session.run( - query, - source=relationship.source, - target=relationship.target, - confidence=relationship.confidence, - ) + # ========================================================== + # Storage + # ========================================================== + @override async def add_fact( self, fact: Fact, ) -> None: assert self.driver is not None - query = cast( + predicate = normalize_predicate( + fact.predicate, + ) + + cypher = cast( LiteralString, f""" MERGE (a:Entity {{name: $subject}}) MERGE (b:Entity {{name: $object}}) - MERGE (a)-[r:{fact.predicate}]->(b) + MERGE (a)-[r:{predicate}]->(b) SET r.confidence = $confidence """, ) async with self.driver.session() as session: await session.run( - query, - # Fixed: subject instead of source + cypher, subject=fact.subject, object=fact.object, confidence=fact.confidence, ) + @override + async def add_facts( + self, + facts: list[Fact], + ) -> None: + for fact in facts: + await self.add_fact(fact) + + @override async def remove_fact( self, fact: Fact, ) -> None: assert self.driver is not None - query = cast( + predicate = normalize_predicate( + fact.predicate, + ) + + cypher = cast( LiteralString, f""" MATCH (a:Entity {{name: $subject}}) - -[r:{fact.predicate}]-> + -[r:{predicate}]-> (b:Entity {{name: $object}}) DELETE r """, @@ -122,23 +165,35 @@ async def remove_fact( async with self.driver.session() as session: await session.run( - query, + cypher, subject=fact.subject, object=fact.object, ) - async def query( + @override + async def remove_facts( + self, + facts: list[Fact], + ) -> None: + for fact in facts: + await self.remove_fact(fact) + + # ========================================================== + # Retrieval + # ========================================================== + + @override + async def search_facts( self, query: str, ) -> list[Fact]: - """ - Return all facts involving the given entity. - """ assert self.driver is not None cypher = """ MATCH (a:Entity)-[r]->(b:Entity) - WHERE a.name = $query OR b.name = $query + WHERE + a.name = $query + OR b.name = $query RETURN a.name AS subject, type(r) AS predicate, @@ -148,7 +203,7 @@ async def query( async with self.driver.session() as session: result = await session.run( - query=cypher, + cypher, parameters={ "query": query, }, @@ -168,17 +223,18 @@ async def query( return facts + @override async def related_entities( self, entity: str, *, depth: int = 1, ) -> list[Entity]: - """ - Return entities connected to the supplied entity. - """ assert self.driver is not None + if depth < 1: + return [] + cypher = cast( LiteralString, f""" @@ -194,10 +250,8 @@ async def related_entities( async with self.driver.session() as session: result = await session.run( - query=cypher, - parameters={ - "entity": entity, - }, + cypher, + entity=entity, ) entities: list[Entity] = [] @@ -206,25 +260,71 @@ async def related_entities( entities.append( Entity( name=record["name"], - label=record["label"], + label=record["label"] or "Entity", ) ) return entities + # ========================================================== + # Schema + # ========================================================== + + @override + async def get_schema(self) -> GraphSchema: + assert self.driver is not None + + async with self.driver.session() as session: + labels_result = await session.run( + """ + CALL db.labels() + YIELD label + RETURN label + ORDER BY label + """ + ) + + labels: list[str] = [] + + async for record in labels_result: + labels.append(record["label"]) + + relationship_result = await session.run( + """ + CALL db.relationshipTypes() + YIELD relationshipType + RETURN relationshipType + ORDER BY relationshipType + """ + ) + + relationship_types: list[str] = [] + + async for record in relationship_result: + relationship_types.append( + record["relationshipType"] + ) + + return GraphSchema( + labels=labels, + relationship_types=relationship_types, + ) + + # ========================================================== + # Maintenance + # ========================================================== + + @override async def clear(self) -> None: - """ - Remove every node and relationship. - """ assert self.driver is not None async with self.driver.session() as session: - await session.run("MATCH (n) DETACH DELETE n") + await session.run( + "MATCH (n) DETACH DELETE n" + ) + @override async def count(self) -> int: - """ - Return the total number of stored relationships. - """ assert self.driver is not None async with self.driver.session() as session: diff --git a/runtime/src/orion/memory/session.py b/runtime/src/orion/memory/session.py index b2b96a3..3eee110 100644 --- a/runtime/src/orion/memory/session.py +++ b/runtime/src/orion/memory/session.py @@ -179,7 +179,7 @@ async def search_facts( ) ) - facts = await self.module.graph.query(query) + facts = await self.module.graph.search_facts(query) await self.event_bus.publish( GraphQueryCompletedEvent( @@ -211,39 +211,59 @@ async def remember_fact( self, fact: Fact, ) -> None: + await self.remember_facts([fact]) - await self.module.graph.add_fact(fact) + async def remember_facts( + self, + facts: list[Fact], + ) -> None: - await self.event_bus.publish( - GraphFactAddedEvent( - session_id=self.session_id, - correlation_id=self.correlation_id, - source=self.SOURCE, - subject=fact.subject, - predicate=fact.predicate, - object=fact.object, - message="Fact added.", + if not facts: + return + + await self.module.graph.add_facts(facts) + + for fact in facts: + await self.event_bus.publish( + GraphFactAddedEvent( + session_id=self.session_id, + correlation_id=self.correlation_id, + source=self.SOURCE, + subject=fact.subject, + predicate=fact.predicate, + object=fact.object, + message="Fact added.", + ) ) - ) async def forget_fact( self, fact: Fact, ) -> None: + await self.forget_facts([fact]) - await self.module.graph.remove_fact(fact) + async def forget_facts( + self, + facts: list[Fact], + ) -> None: - await self.event_bus.publish( - GraphFactRemovedEvent( - session_id=self.session_id, - correlation_id=self.correlation_id, - source=self.SOURCE, - subject=fact.subject, - predicate=fact.predicate, - object=fact.object, - message="Fact removed.", + if not facts: + return + + await self.module.graph.remove_facts(facts) + + for fact in facts: + await self.event_bus.publish( + GraphFactRemovedEvent( + session_id=self.session_id, + correlation_id=self.correlation_id, + source=self.SOURCE, + subject=fact.subject, + predicate=fact.predicate, + object=fact.object, + message="Fact removed.", + ) ) - ) # ------------------------------------------------------------------ # Combined Retrieval diff --git a/runtime/src/orion/orchestrator/config.py b/runtime/src/orion/orchestrator/config.py index 42125c9..279be82 100644 --- a/runtime/src/orion/orchestrator/config.py +++ b/runtime/src/orion/orchestrator/config.py @@ -2,8 +2,9 @@ from dataclasses import dataclass -from langchain_groq import ChatGroq +from langchain_core.language_models.chat_models import BaseChatModel +from orion.integrations._mcp.manager import MCPManager from orion.memory.module import MemoryModule from orion.transport.bridge import IPCBridge @@ -14,6 +15,7 @@ class OrchestratorConfig: Shared resources owned by the orchestrator. """ - llm: ChatGroq + llm: BaseChatModel memory: MemoryModule bridge: IPCBridge + mcp_manager: MCPManager diff --git a/runtime/src/orion/orchestrator/orchestrator.py b/runtime/src/orion/orchestrator/orchestrator.py index 38c49f3..cec696d 100644 --- a/runtime/src/orion/orchestrator/orchestrator.py +++ b/runtime/src/orion/orchestrator/orchestrator.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing_extensions import override + from orion.bus.event_bus import EventBus from orion.orchestrator.config import OrchestratorConfig from orion.runtime.lifecycle import Lifecycle @@ -36,6 +38,7 @@ def __init__( self._started = False + @override async def startup(self) -> None: """ Initialize the ORION runtime. @@ -47,6 +50,7 @@ async def startup(self) -> None: context = ServiceContext( llm=self.config.llm, memory=self.config.memory, + mcp_manager=self.config.mcp_manager, ) self.runtime_services = setup_runtime_services(context) @@ -71,6 +75,7 @@ async def startup(self) -> None: self._started = True + @override async def shutdown(self) -> None: """ Gracefully shutdown the ORION runtime. diff --git a/runtime/src/orion/runtime/run.py b/runtime/src/orion/runtime/run.py index b96c944..9f7ac8c 100644 --- a/runtime/src/orion/runtime/run.py +++ b/runtime/src/orion/runtime/run.py @@ -6,15 +6,20 @@ """ from __future__ import annotations +import os from dotenv import load_dotenv -from langchain_groq import ChatGroq +from langchain_core.language_models.chat_models import BaseChatModel from rich.console import Console from rich.panel import Panel from rich.table import Table from orion.bus.event_bus import EventBus +from orion.integrations._mcp.config import load_config +from orion.integrations._mcp.manager import MCPManager +from orion.llm.config import LLMConfig +from orion.llm.factory import LLMFactory from orion.memory.config import MemoryConfig from orion.memory.module import MemoryModule from orion.memory.planner.planner import RetrievalPlanner @@ -26,16 +31,22 @@ from orion.transport.server import IPCServer console = Console() +load_dotenv() +MCP_CONFIG = os.getenv("MCP_CONFIG", "mcp.json") +SOCKET_PATH = os.getenv("SOCKET_PATH", "/tmp/orion.sock") -def create_llm() -> ChatGroq: +def create_llm() -> BaseChatModel: """Create the application's primary LLM.""" - return ChatGroq( - model="openai/gpt-oss-120b", - temperature=0, + config = LLMConfig() + + providers = LLMFactory.create( + config, ) + return providers.provider.create() + def print_startup_banner() -> None: """Display Orion startup information.""" @@ -66,8 +77,6 @@ def print_startup_banner() -> None: async def run() -> None: """Run the Orion application.""" - load_dotenv() - runtime = OrionRuntime() llm = create_llm() @@ -75,6 +84,8 @@ async def run() -> None: store = SQLiteEventStore() bus = EventBus(store) + mcp_manager = MCPManager(load_config(MCP_CONFIG)) + memory = MemoryModule( config=MemoryConfig(), planner=RetrievalPlanner(llm=llm), @@ -88,16 +99,18 @@ async def run() -> None: llm=llm, memory=memory, bridge=bridge, + mcp_manager=mcp_manager, ), ) server = IPCServer( - socket_path="/tmp/orion.sock", + socket_path=SOCKET_PATH, session_handler=bridge.serve, ) runtime.register(store) runtime.register(memory) + runtime.register(mcp_manager) runtime.register(orchestrator) runtime.register(server) diff --git a/runtime/src/orion/runtime/runtime.py b/runtime/src/orion/runtime/runtime.py index 5feb455..9a70491 100644 --- a/runtime/src/orion/runtime/runtime.py +++ b/runtime/src/orion/runtime/runtime.py @@ -9,6 +9,8 @@ from __future__ import annotations +from typing_extensions import override + from orion.runtime.lifecycle import Lifecycle @@ -42,6 +44,7 @@ def register( self._lifecycles.append(lifecycle) + @override async def startup(self) -> None: """ Start the runtime. @@ -54,6 +57,7 @@ async def startup(self) -> None: self._started = True + @override async def shutdown(self) -> None: """ Shutdown the runtime. diff --git a/runtime/src/orion/services/agent.py b/runtime/src/orion/services/agent.py index 8d1c9a4..59a2ac0 100644 --- a/runtime/src/orion/services/agent.py +++ b/runtime/src/orion/services/agent.py @@ -1,9 +1,11 @@ from __future__ import annotations +import traceback + from dotenv import load_dotenv +from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import HumanMessage from langchain_core.tools import BaseTool -from langchain_groq import ChatGroq from typing_extensions import override from orion.agent.graph import OrionGraph @@ -21,12 +23,15 @@ ResponseCompletedEvent, ResponseStartedEvent, ) -from orion.integrations._mcp.langchain_tools import load_mcp_tools +from orion.integrations._mcp.langchain import create_mcp_tools +from orion.integrations._mcp.manager import MCPManager +from orion.llm.utils import message_content_to_text from orion.memory.models import RetrievedContext from orion.memory.module import MemoryModule from orion.memory.planner.planner import RetrievalPlanner from orion.services.base import BaseService + _ = load_dotenv() @@ -45,21 +50,25 @@ class AgentService(BaseService): ChatPipelineStartEvent, ] - def __init__(self, llm: ChatGroq, memory: MemoryModule) -> None: + def __init__( + self, + llm: BaseChatModel, + memory: MemoryModule, + mcp_manager: MCPManager, + ) -> None: super().__init__() self.llm = llm self.memory = memory + self.mcp_manager = mcp_manager self.planner = RetrievalPlanner( llm=self.llm, ) - self._mcp_tools: list[BaseTool] = [] - @override async def startup(self) -> None: - self._mcp_tools = await load_mcp_tools("mcp.json") + pass @override async def shutdown(self) -> None: @@ -82,18 +91,24 @@ async def handle( ) ) - session = self.memory.session( - session_id=event.session_id, - correlation_id=event.correlation_id, - ) + session_id, correlation_id = self.validate_event(event) - builtin_tools = OrionTools() + mcp_tools = create_mcp_tools( + self.mcp_manager, + session_id=session_id, + correlation_id=correlation_id, + ) tools: list[BaseTool] = [ - *builtin_tools.get_tools(), - *self._mcp_tools, + *OrionTools().get_tools(), + *mcp_tools, ] + session = self.memory.session( + session_id=session_id, + correlation_id=correlation_id, + ) + graph = OrionGraph( retrieve=RetrieveNode(session), agent=AgentNode( @@ -109,10 +124,12 @@ async def handle( ) state: OrionState = { - "session_id": event.session_id, - "correlation_id": event.correlation_id, + "session_id": session_id, + "correlation_id": correlation_id, "messages": [ - HumanMessage(content=event.text), + HumanMessage( + content=event.text, + ), ], "context": RetrievedContext(), } @@ -121,7 +138,15 @@ async def handle( result = await graph.ainvoke(state) response = result["messages"][-1] - assert isinstance(response.content, str) + + response_content = message_content_to_text( + response.content, + ) + + if not response_content: + raise RuntimeError( + "LLM returned an empty response." + ) await self.publish( ResponseStartedEvent( @@ -141,7 +166,7 @@ async def handle( session_id=event.session_id, source=self.service_name, message="Response chunk generated.", - text=response.content, + text=response_content, ) ) @@ -150,18 +175,23 @@ async def handle( correlation_id=event.correlation_id, session_id=event.session_id, source=self.service_name, - text=response.content, + text=response_content, message="Response generation completed.", ) ) except Exception as exc: + error = ( + f"{type(exc).__name__}: {exc}\n" + f"{traceback.format_exc()}" + ) + await self.publish( PipelineFailedEvent( correlation_id=event.correlation_id, session_id=event.session_id, source=self.service_name, message="Agent processing failed.", - error=str(exc), + error=error, ) ) diff --git a/runtime/src/orion/services/base.py b/runtime/src/orion/services/base.py index 05346e2..bd5e15d 100644 --- a/runtime/src/orion/services/base.py +++ b/runtime/src/orion/services/base.py @@ -1,4 +1,6 @@ from abc import ABC, abstractmethod +from uuid import UUID + from orion.bus.event_bus import EventBus from orion.events.base import Event @@ -7,48 +9,66 @@ class BaseService(ABC): """ Base class for all services. - Reponsibilities: - - Provides access to the Event Bus - - Provides event publishing helper - - Provides lifecycle hooks + Responsibilities: + - Provides access to the Event Bus. + - Provides event publishing helper. + - Provides event validation. + - Provides lifecycle hooks. """ service_name: str = "base" subscribed_events: list[type[Event]] = [] - def __init__(self): + def __init__(self) -> None: self.bus = EventBus() async def publish(self, event: Event) -> None: + """Publish an event through the Event Bus.""" await self.bus.publish(event) - @abstractmethod - async def handle(self, event: Event) -> None: + def validate_event( + self, + event: Event, + ) -> tuple[UUID, UUID]: """ - Handle an incoming event + Validate and return the session and correlation identifiers. + + Request-scoped events must contain both identifiers. """ + + assert isinstance( + event.session_id, + UUID, + ), "Event session_id must be a UUID" + + assert isinstance( + event.correlation_id, + UUID, + ), "Event correlation_id must be a UUID" + + return event.session_id, event.correlation_id + + @abstractmethod + async def handle(self, event: Event) -> None: + """Handle an incoming event.""" pass async def startup(self) -> None: - """ - Called during application startup. - Override if needed. - """ + """Called during application startup.""" pass async def shutdown(self) -> None: - """ - Called during application shutdown. - Override if needed. - """ + """Called during application shutdown.""" pass def register(self) -> None: - """ - Register the Service - """ + """Register the service for its subscribed events.""" + for event_type in self.subscribed_events: - self.bus.subscribe(event_type, self.handle) + self.bus.subscribe( + event_type, + self.handle, + ) def __str__(self) -> str: - return f"{self.__class__.__name__}" + return self.__class__.__name__ diff --git a/runtime/src/orion/services/ipc_publisher.py b/runtime/src/orion/services/ipc_publisher.py index c60c638..d2515bc 100644 --- a/runtime/src/orion/services/ipc_publisher.py +++ b/runtime/src/orion/services/ipc_publisher.py @@ -24,6 +24,7 @@ async def handle(self, event: Event) -> None: """ Publish a runtime event over IPC. """ + session_id, correlation_id = self.validate_event(event); if event.type is None: return @@ -45,6 +46,6 @@ async def handle(self, event: Event) -> None: ) await self._bridge.send( - session_id=event.session_id, + session_id=session_id, envelope=envelope, ) diff --git a/runtime/src/orion/services/setup.py b/runtime/src/orion/services/setup.py index de1f181..99976d4 100644 --- a/runtime/src/orion/services/setup.py +++ b/runtime/src/orion/services/setup.py @@ -1,29 +1,32 @@ from dataclasses import dataclass -from langchain_groq import ChatGroq +from langchain_core.language_models.chat_models import BaseChatModel +from orion.integrations._mcp.manager import MCPManager from orion.memory.module import MemoryModule from orion.services.agent import AgentService from orion.services.base import BaseService from orion.services.text_to_speech import TTSService -from orion.services.transcript_generation import TranscriptGenerationService +# from orion.services.transcript_generation import TranscriptGenerationService from orion.services.voice_recording import VoiceRecordingService @dataclass(slots=True) class ServiceContext: - llm: ChatGroq + llm: BaseChatModel memory: MemoryModule + mcp_manager: MCPManager def setup_runtime_services(ctx: ServiceContext) -> list[BaseService]: services: list[BaseService] = [ VoiceRecordingService(), - TranscriptGenerationService(), + # TranscriptGenerationService(), AgentService( llm=ctx.llm, memory=ctx.memory, + mcp_manager=ctx.mcp_manager ), TTSService(), ] diff --git a/runtime/src/orion/services/transcript_generation.py b/runtime/src/orion/services/transcript_generation.py index e3d1a1d..6ddebc6 100644 --- a/runtime/src/orion/services/transcript_generation.py +++ b/runtime/src/orion/services/transcript_generation.py @@ -1,149 +1,149 @@ -import asyncio -import os - -from groq import Groq - -from orion.services.base import BaseService - -from orion.events.base import Event -from orion.events.events import ( - VoiceRecordingCompletedEvent, - TranscriptGeneratedEvent, - TranscriptGenerationFailedEvent, - SilenceDetectedEvent, -) - - -# Phrases Whisper commonly hallucinates from near-silence / ambient noise. -# Compared after lowercasing and stripping surrounding punctuation/space. -_JUNK_PHRASES = { - "", - "you", - "thank you", - "thanks", - "thanks for watching", - "thank you for watching", - "please subscribe", - "subscribe", - "bye", - "goodbye", - "good luck with you", - "okay", - "ok", - "yeah", - "uh", - "um", - "hmm", - "so", - "the", - "you're welcome", - "i'm sorry", -} - -# Transcripts shorter than this (after normalization) are treated as noise. -_MIN_TRANSCRIPT_CHARS = 2 - - -def is_junk_transcript(text: str) -> bool: - """ - True when a transcript is almost certainly a Whisper silence - hallucination rather than real speech, so we can skip the - (expensive) agent + memory pipeline entirely. - """ - normalized = text.strip().lower().strip(" .!?,…\"'") - - if len(normalized) < _MIN_TRANSCRIPT_CHARS: - return True - - return normalized in _JUNK_PHRASES - - -class TranscriptGenerationService(BaseService): - service_name = "transcript_generation" - - subscribed_events = [ - VoiceRecordingCompletedEvent, - ] - - def __init__(self): - super().__init__() - - self.client = Groq( - api_key=os.getenv("GROQ_API_KEY"), - ) - - async def handle( - self, - event: Event, - ) -> None: - - try: - assert isinstance( - event, - VoiceRecordingCompletedEvent, - ) - - transcript = await self.transcribe( - str(event.audio_path), - ) - - # - # Guard: Whisper hallucinates phantom phrases ("Thank you.", - # ".", "you") on near-silence. Running the agent + memory - # pipeline on those wastes a large number of LLM tokens, so - # we drop them here and let the loop simply listen again. - # - if is_junk_transcript(transcript): - await self.publish( - SilenceDetectedEvent( - session_id=event.session_id, - correlation_id=event.correlation_id, - source=self.service_name, - message="Ignored non-speech transcript", - ) - ) - return - - await self.publish( - TranscriptGeneratedEvent( - session_id=event.session_id, - correlation_id=event.correlation_id, - source=self.service_name, - message="Transcript generated", - text=transcript, - ) - ) - - except Exception as e: - await self.publish( - TranscriptGenerationFailedEvent( - session_id=event.session_id, - correlation_id=event.correlation_id, - source=self.service_name, - message="Transcript generation failed", - error=str(e), - ) - ) - - async def transcribe( - self, - audio_path: str, - ) -> str: - - return await asyncio.to_thread( - self._transcribe_sync, - audio_path, - ) - - def _transcribe_sync( - self, - audio_path: str, - ) -> str: - - with open(audio_path, "rb") as audio_file: - transcription = self.client.audio.transcriptions.create( - file=audio_file, - model="whisper-large-v3-turbo", - ) - - return transcription.text +# import asyncio +# import os + +# from groq import Groq + +# from orion.services.base import BaseService + +# from orion.events.base import Event +# from orion.events.events import ( +# VoiceRecordingCompletedEvent, +# TranscriptGeneratedEvent, +# TranscriptGenerationFailedEvent, +# SilenceDetectedEvent, +# ) + + +# # Phrases Whisper commonly hallucinates from near-silence / ambient noise. +# # Compared after lowercasing and stripping surrounding punctuation/space. +# _JUNK_PHRASES = { +# "", +# "you", +# "thank you", +# "thanks", +# "thanks for watching", +# "thank you for watching", +# "please subscribe", +# "subscribe", +# "bye", +# "goodbye", +# "good luck with you", +# "okay", +# "ok", +# "yeah", +# "uh", +# "um", +# "hmm", +# "so", +# "the", +# "you're welcome", +# "i'm sorry", +# } + +# # Transcripts shorter than this (after normalization) are treated as noise. +# _MIN_TRANSCRIPT_CHARS = 2 + + +# def is_junk_transcript(text: str) -> bool: +# """ +# True when a transcript is almost certainly a Whisper silence +# hallucination rather than real speech, so we can skip the +# (expensive) agent + memory pipeline entirely. +# """ +# normalized = text.strip().lower().strip(" .!?,…\"'") + +# if len(normalized) < _MIN_TRANSCRIPT_CHARS: +# return True + +# return normalized in _JUNK_PHRASES + + +# class TranscriptGenerationService(BaseService): +# service_name = "transcript_generation" + +# subscribed_events = [ +# VoiceRecordingCompletedEvent, +# ] + +# def __init__(self): +# super().__init__() + +# self.client = Groq( +# api_key=os.getenv("GROQ_API_KEY"), +# ) + +# async def handle( +# self, +# event: Event, +# ) -> None: + +# try: +# assert isinstance( +# event, +# VoiceRecordingCompletedEvent, +# ) + +# transcript = await self.transcribe( +# str(event.audio_path), +# ) + +# # +# # Guard: Whisper hallucinates phantom phrases ("Thank you.", +# # ".", "you") on near-silence. Running the agent + memory +# # pipeline on those wastes a large number of LLM tokens, so +# # we drop them here and let the loop simply listen again. +# # +# if is_junk_transcript(transcript): +# await self.publish( +# SilenceDetectedEvent( +# session_id=event.session_id, +# correlation_id=event.correlation_id, +# source=self.service_name, +# message="Ignored non-speech transcript", +# ) +# ) +# return + +# await self.publish( +# TranscriptGeneratedEvent( +# session_id=event.session_id, +# correlation_id=event.correlation_id, +# source=self.service_name, +# message="Transcript generated", +# text=transcript, +# ) +# ) + +# except Exception as e: +# await self.publish( +# TranscriptGenerationFailedEvent( +# session_id=event.session_id, +# correlation_id=event.correlation_id, +# source=self.service_name, +# message="Transcript generation failed", +# error=str(e), +# ) +# ) + +# async def transcribe( +# self, +# audio_path: str, +# ) -> str: + +# return await asyncio.to_thread( +# self._transcribe_sync, +# audio_path, +# ) + +# def _transcribe_sync( +# self, +# audio_path: str, +# ) -> str: + +# with open(audio_path, "rb") as audio_file: +# transcription = self.client.audio.transcriptions.create( +# file=audio_file, +# model="whisper-large-v3-turbo", +# ) + +# return transcription.text diff --git a/runtime/src/orion/transport/bridge.py b/runtime/src/orion/transport/bridge.py index 1a6002f..a1a3b5e 100644 --- a/runtime/src/orion/transport/bridge.py +++ b/runtime/src/orion/transport/bridge.py @@ -67,6 +67,8 @@ async def send( """ Send an IPC message to a connected client. """ + assert isinstance(session_id, UUID) + session = self._sessions.get(session_id) if session is None: diff --git a/runtime/src/orion/transport/messages.py b/runtime/src/orion/transport/messages.py index 7d3570e..77b5696 100644 --- a/runtime/src/orion/transport/messages.py +++ b/runtime/src/orion/transport/messages.py @@ -7,13 +7,18 @@ from pydantic import BaseModel, ConfigDict, Field +# ====================================================================== +# Message Types +# ====================================================================== + + class MessageType(StrEnum): """ All message types supported by the Orion IPC protocol. - These messages define the contract between a client (TUI, mobile, - web, etc.) and the Orion runtime. Transport messages are intentionally - independent of the runtime's internal domain events. + Transport messages are intentionally independent of Orion's internal + domain event classes. Internal events are translated into Envelopes + before being sent to clients. """ # ------------------------------------------------------------------ @@ -72,6 +77,22 @@ class MessageType(StrEnum): #: Indicates that a tool execution has completed. TOOL_FINISHED = "tool_finished" + # ------------------------------------------------------------------ + # Pipeline + # ------------------------------------------------------------------ + + #: Indicates that a processing pipeline has started. + PIPELINE_STARTED = "pipeline_started" + + #: Indicates that a processing pipeline completed successfully. + PIPELINE_COMPLETED = "pipeline_completed" + + #: Indicates that a processing pipeline failed. + PIPELINE_FAILED = "pipeline_failed" + + #: Indicates that a processing pipeline was restarted. + PIPELINE_RESTARTED = "pipeline_restarted" + # ------------------------------------------------------------------ # Runtime # ------------------------------------------------------------------ @@ -79,28 +100,41 @@ class MessageType(StrEnum): #: General runtime status update. STATUS = "status" - #: Indicates an unrecoverable error. + #: Indicates an error. ERROR = "error" + #: Runtime event + RUNTIME_EVENT = "runtime_event" + + +# ====================================================================== +# Envelope +# ====================================================================== + class Envelope(BaseModel): """ - A transport message exchanged between the Orion runtime and a client. + Transport envelope exchanged between the Orion runtime and clients. + + Every message transmitted over IPC is wrapped inside an Envelope. - Every message transmitted over the IPC channel is wrapped inside an - Envelope. The payload is interpreted according to the message type. + The payload is interpreted according to the message type. """ - model_config = ConfigDict(use_enum_values=True) + model_config = ConfigDict( + use_enum_values=True, + ) #: Protocol version used for compatibility checks. version: int = 1 - #: Unique identifier for this message. - id: UUID = Field(default_factory=uuid4) + #: Unique identifier for this individual message. + id: UUID = Field( + default_factory=uuid4, + ) - #: ID of the request/workflow this message belongs to - correlation_id: UUID = Field(default_factory=uuid4) + #: Identifier of the request/workflow this message belongs to. + correlation_id: UUID | None = None #: Type of message being transmitted. type: MessageType @@ -152,13 +186,13 @@ class VoiceChunkPayload(BaseModel): class VoiceEndPayload(BaseModel): - """Marks the end of a voice recording. + """ + Marks the end of a voice recording. - Carries the path to the file the client recorded; the runtime reads and - transcribes it (the preferred, non-streamed flow). + Carries the path to the file the runtime should process. """ - #: Filesystem path to the recorded audio, as seen by the runtime. + #: Filesystem path to the recorded audio. path: str @@ -167,6 +201,11 @@ class VoiceEndPayload(BaseModel): # ====================================================================== +# ---------------------------------------------------------------------- +# Assistant +# ---------------------------------------------------------------------- + + class AssistantStartPayload(BaseModel): """Signals the start of assistant response generation.""" @@ -182,6 +221,59 @@ class AssistantEndPayload(BaseModel): """Signals completion of assistant response generation.""" +# ---------------------------------------------------------------------- +# Pipeline +# ---------------------------------------------------------------------- + + +class PipelineStartedPayload(BaseModel): + """Signals that a processing pipeline has started.""" + + #: Name of the pipeline. + pipeline: str + + +class PipelineCompletedPayload(BaseModel): + """Signals that a processing pipeline completed successfully.""" + + #: Name of the pipeline. + pipeline: str + + +class PipelineFailedPayload(BaseModel): + """ + Describes a failed processing pipeline. + + The traceback is optional so production clients do not have to + expose implementation details while development clients can still + display complete debugging information. + """ + + #: Name of the pipeline that failed. + pipeline: str + + #: Machine-readable exception type. + error_type: str + + #: Human-readable error message. + message: str + + #: Full traceback when available. + traceback: str | None = None + + +class PipelineRestartedPayload(BaseModel): + """Signals that a processing pipeline has been restarted.""" + + #: Name of the pipeline. + pipeline: str + + +# ---------------------------------------------------------------------- +# Tool Execution +# ---------------------------------------------------------------------- + + class ToolStartedPayload(BaseModel): """Indicates that the assistant has started executing a tool.""" @@ -199,6 +291,11 @@ class ToolFinishedPayload(BaseModel): success: bool +# ---------------------------------------------------------------------- +# Runtime Status +# ---------------------------------------------------------------------- + + class StatusPayload(BaseModel): """General runtime status update.""" @@ -207,7 +304,12 @@ class StatusPayload(BaseModel): class ErrorPayload(BaseModel): - """Represents an error returned by the runtime.""" + """ + Represents an error returned by the runtime. + + This is intended for errors that are not represented by a more + specific pipeline failure message. + """ #: Machine-readable error identifier. code: str @@ -215,10 +317,70 @@ class ErrorPayload(BaseModel): #: Human-readable error description. message: str + #: Exception type when available. + error_type: str | None = None + + #: Pipeline associated with the error when applicable. + pipeline: str | None = None + + #: Full traceback when available. + traceback: str | None = None + + +# ---------------------------------------------------------------------- +# Voice / Speech +# ---------------------------------------------------------------------- + + +class VoiceRecordingStartedPayload(BaseModel): + """Signals that voice recording has started.""" + + +class VoiceRecordingCompletedPayload(BaseModel): + """Signals that voice recording has completed.""" + + #: Path to the recorded audio. + path: str | None = None + + +class TranscriptPayload(BaseModel): + """Contains a generated speech transcript.""" + + #: Transcribed text. + text: str + + +class SpeechSynthesisPayload(BaseModel): + """Contains information about generated speech.""" + + #: Text that was synthesized. + text: str + + #: Generated audio path, if available. + audio_path: str | None = None + + +# ---------------------------------------------------------------------- +# Audio Playback +# ---------------------------------------------------------------------- + + +class AudioPlaybackStartedPayload(BaseModel): + """Signals that audio playback has started.""" + + +class AudioPlaybackCompletedPayload(BaseModel): + """Signals that audio playback has completed.""" + + +# ---------------------------------------------------------------------- +# Connection +# ---------------------------------------------------------------------- + class PingPayload(BaseModel): """Ping request payload.""" class PongPayload(BaseModel): - """Ping response payload.""" + """Ping response payload.""" \ No newline at end of file diff --git a/runtime/tests/integrations/test_mcp_discovery.py b/runtime/tests/integrations/test_mcp_discovery.py new file mode 100644 index 0000000..012c5d1 --- /dev/null +++ b/runtime/tests/integrations/test_mcp_discovery.py @@ -0,0 +1,183 @@ +# runtime/tests/integrations/test_mcp_discovery.py + +from __future__ import annotations + +from types import SimpleNamespace + +from orion.integrations._mcp.discovery import ( + mcp_tools_to_openai, +) + + +def test_converts_single_tool() -> None: + tool = SimpleNamespace( + name="read_file", + description="Read a file", + inputSchema={ + "type": "object", + "properties": { + "path": { + "type": "string", + }, + }, + "required": ["path"], + }, + ) + + result = SimpleNamespace( + tools=[tool], + ) + + schemas = mcp_tools_to_openai(result) + + assert schemas == [ + { + "type": "function", + "function": { + "name": "read_file", + "description": "Read a file", + "parameters": { + "type": "object", + "properties": { + "path": { + "type": "string", + }, + }, + "required": ["path"], + }, + }, + } + ] + + +def test_converts_multiple_tools() -> None: + result = SimpleNamespace( + tools=[ + SimpleNamespace( + name="read_file", + description="Read a file", + inputSchema={ + "type": "object", + "properties": { + "path": { + "type": "string", + } + }, + }, + ), + SimpleNamespace( + name="list_directory", + description="List directory contents", + inputSchema={ + "type": "object", + "properties": { + "path": { + "type": "string", + } + }, + }, + ), + ] + ) + + schemas = mcp_tools_to_openai(result) + + assert len(schemas) == 2 + + assert schemas[0]["function"]["name"] == "read_file" + assert ( + schemas[0]["function"]["description"] + == "Read a file" + ) + + assert schemas[1]["function"]["name"] == "list_directory" + assert ( + schemas[1]["function"]["description"] + == "List directory contents" + ) + + +def test_empty_tools_returns_empty_list() -> None: + result = SimpleNamespace( + tools=[] + ) + + schemas = mcp_tools_to_openai(result) + + assert schemas == [] + + +def test_missing_description_becomes_empty_string() -> None: + tool = SimpleNamespace( + name="test_tool", + description=None, + inputSchema={ + "type": "object", + "properties": {}, + }, + ) + + result = SimpleNamespace( + tools=[tool] + ) + + schemas = mcp_tools_to_openai(result) + + assert ( + schemas[0]["function"]["description"] + == "" + ) + + +def test_empty_input_schema_gets_default_schema() -> None: + tool = SimpleNamespace( + name="test_tool", + description="Test tool", + inputSchema=None, + ) + + result = SimpleNamespace( + tools=[tool] + ) + + schemas = mcp_tools_to_openai(result) + + assert ( + schemas[0]["function"]["parameters"] + == { + "type": "object", + "properties": {}, + } + ) + + +def test_preserves_input_schema() -> None: + input_schema = { + "type": "object", + "properties": { + "query": { + "type": "string", + }, + "limit": { + "type": "integer", + }, + }, + "required": ["query"], + } + + tool = SimpleNamespace( + name="search", + description="Search", + inputSchema=input_schema, + ) + + result = SimpleNamespace( + tools=[tool] + ) + + schemas = mcp_tools_to_openai(result) + + assert ( + schemas[0]["function"]["parameters"] + is input_schema + ) diff --git a/runtime/tests/integrations/test_mcp_langchain.py b/runtime/tests/integrations/test_mcp_langchain.py new file mode 100644 index 0000000..dcddf4a --- /dev/null +++ b/runtime/tests/integrations/test_mcp_langchain.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest + +from orion.integrations._mcp.langchain import create_mcp_tools + + +# ========================================================== +# Helpers +# ========================================================== + + +def make_tool_schema( + name: str, + description: str = "", + parameters: dict | None = None, +) -> dict: + return { + "type": "function", + "function": { + "name": name, + "description": description, + "parameters": parameters + or { + "type": "object", + "properties": {}, + }, + }, + } + + +# ========================================================== +# Fixtures +# ========================================================== + + +@pytest.fixture +def session_id(): + return uuid4() + + +@pytest.fixture +def correlation_id(): + return uuid4() + + +# ========================================================== +# Empty Manager +# ========================================================== + + +def test_create_mcp_tools_empty( + session_id, + correlation_id, +) -> None: + manager = AsyncMock() + manager.tools = [] + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + assert tools == [] + + +# ========================================================== +# Tool Conversion +# ========================================================== + + +def test_create_mcp_tools( + session_id, + correlation_id, +) -> None: + manager = AsyncMock() + + manager.tools = [ + make_tool_schema( + name="read_file", + description="Read a file.", + ) + ] + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + assert len(tools) == 1 + + tool = tools[0] + + assert tool.name == "read_file" + assert tool.description == "Read a file." + + +def test_create_multiple_mcp_tools( + session_id, + correlation_id, +) -> None: + manager = AsyncMock() + + manager.tools = [ + make_tool_schema( + name="read_file", + description="Read a file.", + ), + make_tool_schema( + name="write_file", + description="Write a file.", + ), + make_tool_schema( + name="list_directory", + description="List a directory.", + ), + ] + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + assert len(tools) == 3 + + assert [tool.name for tool in tools] == [ + "read_file", + "write_file", + "list_directory", + ] + + +# ========================================================== +# Schema +# ========================================================== + + +def test_create_mcp_tool_preserves_schema( + session_id, + correlation_id, +) -> None: + manager = AsyncMock() + + parameters = { + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path to the file.", + }, + }, + "required": ["path"], + } + + manager.tools = [ + make_tool_schema( + name="read_file", + description="Read a file.", + parameters=parameters, + ) + ] + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + tool = tools[0] + + assert tool.args == parameters + + +def test_create_mcp_tool_uses_default_schema( + session_id, + correlation_id, +) -> None: + manager = AsyncMock() + + manager.tools = [ + { + "type": "function", + "function": { + "name": "ping", + "description": "Ping the server.", + }, + } + ] + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + tool = tools[0] + + assert tool.args == { + "properties": {}, + "type": "object", + } + + +# ========================================================== +# Tool Execution +# ========================================================== + + +@pytest.mark.asyncio +async def test_mcp_tool_calls_manager( + session_id, + correlation_id, +) -> None: + manager = AsyncMock() + + manager.tools = [ + make_tool_schema( + name="read_file", + description="Read a file.", + ) + ] + + manager.call_tool.return_value = "file contents" + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + result = await tools[0].ainvoke( + { + "path": "/tmp/test.txt", + } + ) + + assert result == "file contents" + + manager.call_tool.assert_awaited_once_with( + "read_file", + { + "path": "/tmp/test.txt", + }, + session_id=session_id, + correlation_id=correlation_id, + ) + + +@pytest.mark.asyncio +async def test_mcp_tools_call_correct_tool( + session_id, + correlation_id, +) -> None: + """ + Ensure every generated tool captures its own MCP tool name. + """ + + manager = AsyncMock() + + manager.tools = [ + make_tool_schema("read_file"), + make_tool_schema("write_file"), + make_tool_schema("list_directory"), + ] + + manager.call_tool.side_effect = [ + "read result", + "write result", + "list result", + ] + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + assert await tools[0].ainvoke({"path": "a"}) == "read result" + assert await tools[1].ainvoke({"path": "b"}) == "write result" + assert await tools[2].ainvoke({"path": "c"}) == "list result" + + assert manager.call_tool.await_args_list[0].args == ( + "read_file", + {"path": "a"}, + ) + + assert manager.call_tool.await_args_list[1].args == ( + "write_file", + {"path": "b"}, + ) + + assert manager.call_tool.await_args_list[2].args == ( + "list_directory", + {"path": "c"}, + ) + + +# ========================================================== +# Request Context +# ========================================================== + + +@pytest.mark.asyncio +async def test_mcp_tool_passes_request_context( + session_id, + correlation_id, +) -> None: + manager = AsyncMock() + + manager.tools = [ + make_tool_schema("read_file"), + ] + + manager.call_tool.return_value = "result" + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + await tools[0].ainvoke( + { + "path": "/tmp/test.txt", + } + ) + + manager.call_tool.assert_awaited_once_with( + "read_file", + { + "path": "/tmp/test.txt", + }, + session_id=session_id, + correlation_id=correlation_id, + ) + + +# ========================================================== +# Error Propagation +# ========================================================== + + +@pytest.mark.asyncio +async def test_mcp_tool_propagates_manager_failure( + session_id, + correlation_id, +) -> None: + manager = AsyncMock() + + manager.tools = [ + make_tool_schema("read_file"), + ] + + manager.call_tool.side_effect = RuntimeError( + "MCP tool execution failed" + ) + + tools = create_mcp_tools( + manager, + session_id=session_id, + correlation_id=correlation_id, + ) + + with pytest.raises( + RuntimeError, + match="MCP tool execution failed", + ): + await tools[0].ainvoke( + { + "path": "/tmp/test.txt", + } + ) diff --git a/runtime/tests/integrations/test_mcp_manager.py b/runtime/tests/integrations/test_mcp_manager.py new file mode 100644 index 0000000..e7d1338 --- /dev/null +++ b/runtime/tests/integrations/test_mcp_manager.py @@ -0,0 +1,806 @@ +# runtime/tests/integrations/test_mcp_manager.py + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 + +import pytest + +from orion.bus.event_bus import EventBus +from orion.core.singleton import SingletonMeta +from orion.integrations._mcp.config import MCPConfig, MCPServerConfig +from orion.integrations._mcp.events import ( + MCPServerShutdownFailedEvent, + MCPServerStartedEvent, + MCPServerStartupFailedEvent, + MCPServerStoppedEvent, + MCPToolCalledEvent, + MCPToolCompletedEvent, + MCPToolFailedEvent, + MCPToolsDiscoveredEvent, + MCPToolsDiscoveryFailedEvent, +) +from orion.integrations._mcp.manager import MCPManager +from orion.integrations._mcp.server import MCPServer + + +# ========================================================== +# Fixtures +# ========================================================== + +@pytest.fixture(autouse=True) +def reset_event_bus(): + """ + Reset the EventBus singleton between tests. + + EventBus is process-wide and therefore must not leak + subscribers or stores between tests. + """ + + SingletonMeta._instances.pop(EventBus, None) + + yield + + SingletonMeta._instances.pop(EventBus, None) + + +@pytest.fixture +def event_bus() -> EventBus: + """ + Create the global EventBus with a mocked event store. + """ + + store = MagicMock() + store.append = AsyncMock() + + return EventBus(store) + + +@pytest.fixture +def config() -> MCPConfig: + return MCPConfig( + servers=[ + MCPServerConfig( + name="filesystem", + transport="stdio", + command="npx", + ), + ] + ) + + +@pytest.fixture +def manager( + config: MCPConfig, + event_bus: EventBus, +) -> MCPManager: + """ + EventBus must exist before MCPManager is constructed. + """ + + return MCPManager(config) + + +# ========================================================== +# Helpers +# ========================================================== + + +async def collect_events( + event_bus: EventBus, +) -> list[object]: + """ + Subscribe to every event and return the collected list. + """ + + events: list[object] = [] + + async def capture(event: object) -> None: + events.append(event) + + event_bus.subscribe_all(capture) + + return events + + +# ========================================================== +# Properties +# ========================================================== + + +def test_servers_initially_empty( + manager: MCPManager, +) -> None: + assert manager.servers == {} + + +def test_tools_initially_empty( + manager: MCPManager, +) -> None: + assert manager.tools == [] + + +def test_started_initially_false( + manager: MCPManager, +) -> None: + assert manager.started is False + + +# ========================================================== +# Startup +# ========================================================== + + +@pytest.mark.asyncio +async def test_startup_connects_server_and_discovers_tools( + manager: MCPManager, + event_bus: EventBus, +) -> None: + mock_server = MagicMock(spec=MCPServer) + + mock_server.name = "filesystem" + mock_server.startup = AsyncMock() + mock_server.list_tools = AsyncMock( + return_value=MagicMock() + ) + + schemas = [ + { + "type": "function", + "function": { + "name": "read_text_file", + "description": "Read a file", + "parameters": {}, + }, + }, + { + "type": "function", + "function": { + "name": "list_directory", + "description": "List a directory", + "parameters": {}, + }, + }, + ] + + events = await collect_events(event_bus) + + with ( + patch( + "orion.integrations._mcp.manager.MCPServer", + return_value=mock_server, + ), + patch( + "orion.integrations._mcp.manager.mcp_tools_to_openai", + return_value=schemas, + ), + ): + await manager.startup() + + assert manager.started is True + assert manager.servers["filesystem"] is mock_server + + assert len(manager.tools) == 2 + assert ( + manager.tools[0]["function"]["name"] + == "read_text_file" + ) + assert ( + manager.tools[1]["function"]["name"] + == "list_directory" + ) + + mock_server.startup.assert_awaited_once() + mock_server.list_tools.assert_awaited_once() + + started = [ + event + for event in events + if isinstance(event, MCPServerStartedEvent) + ] + + discovered = [ + event + for event in events + if isinstance(event, MCPToolsDiscoveredEvent) + ] + + assert len(started) == 1 + assert started[0].server_name == "filesystem" + assert started[0].transport == "stdio" + + assert len(discovered) == 1 + assert discovered[0].server_name == "filesystem" + assert discovered[0].tool_count == 2 + assert discovered[0].tool_names == [ + "read_text_file", + "list_directory", + ] + + +@pytest.mark.asyncio +async def test_startup_skips_disabled_servers( + event_bus: EventBus, +) -> None: + config = MCPConfig( + servers=[ + MCPServerConfig( + name="disabled", + command="echo", + enabled=False, + ) + ] + ) + + manager = MCPManager(config) + + with patch( + "orion.integrations._mcp.manager.MCPServer" + ) as mock_server: + await manager.startup() + + mock_server.assert_not_called() + + assert manager.started is True + assert manager.servers == {} + assert manager.tools == [] + + +@pytest.mark.asyncio +async def test_startup_failure_does_not_stop_other_servers( + event_bus: EventBus, +) -> None: + config = MCPConfig( + servers=[ + MCPServerConfig( + name="broken", + command="broken", + ), + MCPServerConfig( + name="working", + command="working", + ), + ] + ) + + manager = MCPManager(config) + + broken_server = MagicMock(spec=MCPServer) + broken_server.name = "broken" + broken_server.startup = AsyncMock( + side_effect=RuntimeError("connection failed") + ) + + working_server = MagicMock(spec=MCPServer) + working_server.name = "working" + working_server.startup = AsyncMock() + working_server.list_tools = AsyncMock( + return_value=MagicMock() + ) + + events = await collect_events(event_bus) + + with ( + patch( + "orion.integrations._mcp.manager.MCPServer", + side_effect=[ + broken_server, + working_server, + ], + ), + patch( + "orion.integrations._mcp.manager.mcp_tools_to_openai", + return_value=[], + ), + ): + await manager.startup() + + assert manager.started is True + + assert "broken" not in manager.servers + assert "working" in manager.servers + + failures = [ + event + for event in events + if isinstance( + event, + MCPServerStartupFailedEvent, + ) + ] + + assert len(failures) == 1 + assert failures[0].server_name == "broken" + assert failures[0].transport == "stdio" + assert failures[0].error == "connection failed" + + assert failures[0].correlation_id is None + assert failures[0].session_id is None + + +# ========================================================== +# Tool Discovery Failure +# ========================================================== + + +@pytest.mark.asyncio +async def test_tool_discovery_failure_publishes_event( + manager: MCPManager, + event_bus: EventBus, +) -> None: + mock_server = MagicMock(spec=MCPServer) + + mock_server.name = "filesystem" + mock_server.startup = AsyncMock() + mock_server.list_tools = AsyncMock( + side_effect=RuntimeError("discovery failed") + ) + + events = await collect_events(event_bus) + + with patch( + "orion.integrations._mcp.manager.MCPServer", + return_value=mock_server, + ): + await manager.startup() + + assert "filesystem" in manager.servers + assert manager.tools == [] + + failures = [ + event + for event in events + if isinstance( + event, + MCPToolsDiscoveryFailedEvent, + ) + ] + + assert len(failures) == 1 + assert failures[0].server_name == "filesystem" + assert failures[0].error == "discovery failed" + + # Discovery is a lifecycle operation, not request-scoped. + assert failures[0].correlation_id is None + assert failures[0].session_id is None + + +# ========================================================== +# Duplicate Tools +# ========================================================== + + +@pytest.mark.asyncio +async def test_duplicate_tool_is_rejected( + event_bus: EventBus, +) -> None: + config = MCPConfig( + servers=[ + MCPServerConfig( + name="server_a", + command="a", + ), + MCPServerConfig( + name="server_b", + command="b", + ), + ] + ) + + manager = MCPManager(config) + + server_a = MagicMock(spec=MCPServer) + server_a.name = "server_a" + server_a.startup = AsyncMock() + server_a.list_tools = AsyncMock( + return_value=MagicMock() + ) + + server_b = MagicMock(spec=MCPServer) + server_b.name = "server_b" + server_b.startup = AsyncMock() + server_b.list_tools = AsyncMock( + return_value=MagicMock() + ) + + schemas = [ + { + "type": "function", + "function": { + "name": "same_tool", + "description": "Duplicate", + "parameters": {}, + }, + } + ] + + with ( + patch( + "orion.integrations._mcp.manager.MCPServer", + side_effect=[ + server_a, + server_b, + ], + ), + patch( + "orion.integrations._mcp.manager.mcp_tools_to_openai", + return_value=schemas, + ), + ): + await manager.startup() + + # First server owns the tool. + assert len(manager.tools) == 1 + assert ( + manager.tools[0]["function"]["name"] + == "same_tool" + ) + assert ( + manager._tool_routing["same_tool"] + == "server_a" + ) + + +# ========================================================== +# Startup Idempotency +# ========================================================== + + +@pytest.mark.asyncio +async def test_startup_only_runs_once( + manager: MCPManager, + event_bus: EventBus, +) -> None: + mock_server = MagicMock(spec=MCPServer) + + mock_server.name = "filesystem" + mock_server.startup = AsyncMock() + mock_server.list_tools = AsyncMock( + return_value=MagicMock() + ) + + with ( + patch( + "orion.integrations._mcp.manager.MCPServer", + return_value=mock_server, + ), + patch( + "orion.integrations._mcp.manager.mcp_tools_to_openai", + return_value=[], + ), + ): + await manager.startup() + await manager.startup() + + mock_server.startup.assert_awaited_once() + + +# ========================================================== +# Server Access +# ========================================================== + + +def test_server_returns_connected_server( + manager: MCPManager, +) -> None: + mock_server = MagicMock(spec=MCPServer) + + manager._servers["filesystem"] = mock_server + + assert manager.server("filesystem") is mock_server + + +def test_server_unknown_name_raises( + manager: MCPManager, +) -> None: + with pytest.raises( + ValueError, + match="Unknown MCP server: missing", + ): + manager.server("missing") + + +# ========================================================== +# Raw Tool Execution +# ========================================================== + + +@pytest.mark.asyncio +async def test_call_tool_raw( + manager: MCPManager, + event_bus: EventBus, +) -> None: + mock_server = MagicMock(spec=MCPServer) + + mock_result = MagicMock() + + mock_server.call_tool = AsyncMock( + return_value=mock_result + ) + + manager._servers["filesystem"] = mock_server + manager._tool_routing["read_file"] = "filesystem" + + session_id = uuid4() + correlation_id = uuid4() + + events = await collect_events(event_bus) + + result = await manager.call_tool_raw( + "read_file", + {"path": "test.txt"}, + session_id=session_id, + correlation_id=correlation_id, + ) + + assert result is mock_result + + mock_server.call_tool.assert_awaited_once_with( + "read_file", + {"path": "test.txt"}, + ) + + called = [ + event + for event in events + if isinstance( + event, + MCPToolCalledEvent, + ) + ] + + completed = [ + event + for event in events + if isinstance( + event, + MCPToolCompletedEvent, + ) + ] + + assert len(called) == 1 + assert len(completed) == 1 + + assert called[0].server_name == "filesystem" + assert called[0].tool_name == "read_file" + assert called[0].session_id == session_id + assert called[0].correlation_id == correlation_id + + assert completed[0].server_name == "filesystem" + assert completed[0].tool_name == "read_file" + assert completed[0].session_id == session_id + assert completed[0].correlation_id == correlation_id + + +@pytest.mark.asyncio +async def test_call_tool_raw_unknown_tool( + manager: MCPManager, +) -> None: + with pytest.raises( + ValueError, + match="Unknown MCP tool: missing", + ): + await manager.call_tool_raw( + "missing", + {}, + session_id=uuid4(), + correlation_id=uuid4(), + ) + + +# ========================================================== +# Tool Execution Failure +# ========================================================== + + +@pytest.mark.asyncio +async def test_call_tool_raw_failure_publishes_event( + manager: MCPManager, + event_bus: EventBus, +) -> None: + mock_server = MagicMock(spec=MCPServer) + + mock_server.call_tool = AsyncMock( + side_effect=RuntimeError("tool failed") + ) + + manager._servers["filesystem"] = mock_server + manager._tool_routing["read_file"] = "filesystem" + + session_id = uuid4() + correlation_id = uuid4() + + events = await collect_events(event_bus) + + with pytest.raises( + RuntimeError, + match="tool failed", + ): + await manager.call_tool_raw( + "read_file", + {}, + session_id=session_id, + correlation_id=correlation_id, + ) + + failures = [ + event + for event in events + if isinstance( + event, + MCPToolFailedEvent, + ) + ] + + assert len(failures) == 1 + + assert failures[0].server_name == "filesystem" + assert failures[0].tool_name == "read_file" + assert failures[0].error == "tool failed" + + assert failures[0].session_id == session_id + assert failures[0].correlation_id == correlation_id + + +# ========================================================== +# call_tool() +# ========================================================== + + +@pytest.mark.asyncio +async def test_call_tool_flattens_text( + manager: MCPManager, + event_bus: EventBus, +) -> None: + mock_server = MagicMock(spec=MCPServer) + + text_one = MagicMock() + text_one.text = "first" + + text_two = MagicMock() + text_two.text = "second" + + result = MagicMock() + result.content = [ + text_one, + text_two, + ] + + mock_server.call_tool = AsyncMock( + return_value=result + ) + + manager._servers["filesystem"] = mock_server + manager._tool_routing["read_file"] = "filesystem" + + output = await manager.call_tool( + "read_file", + {}, + session_id=uuid4(), + correlation_id=uuid4(), + ) + + assert output == "first\nsecond" + + +@pytest.mark.asyncio +async def test_call_tool_returns_no_output( + manager: MCPManager, +) -> None: + mock_server = MagicMock(spec=MCPServer) + + result = MagicMock() + result.content = [] + + mock_server.call_tool = AsyncMock( + return_value=result + ) + + manager._servers["filesystem"] = mock_server + manager._tool_routing["empty_tool"] = "filesystem" + + output = await manager.call_tool( + "empty_tool", + {}, + session_id=uuid4(), + correlation_id=uuid4(), + ) + + assert output == "(no output)" + + +# ========================================================== +# Shutdown +# ========================================================== + + +@pytest.mark.asyncio +async def test_shutdown_stops_servers( + manager: MCPManager, + event_bus: EventBus, +) -> None: + server_one = MagicMock(spec=MCPServer) + server_one.name = "one" + server_one.shutdown = AsyncMock() + + server_two = MagicMock(spec=MCPServer) + server_two.name = "two" + server_two.shutdown = AsyncMock() + + manager._servers["one"] = server_one + manager._servers["two"] = server_two + + manager._tool_routing["tool_one"] = "one" + manager._tools.append( + { + "type": "function", + "function": { + "name": "tool_one", + }, + } + ) + + manager._started = True + + events = await collect_events(event_bus) + + await manager.shutdown() + + server_one.shutdown.assert_awaited_once() + server_two.shutdown.assert_awaited_once() + + assert manager.servers == {} + assert manager.tools == [] + assert manager._tool_routing == {} + assert manager.started is False + + stopped = [ + event + for event in events + if isinstance( + event, + MCPServerStoppedEvent, + ) + ] + + assert len(stopped) == 2 + + +@pytest.mark.asyncio +async def test_shutdown_failure_publishes_event( + manager: MCPManager, + event_bus: EventBus, +) -> None: + mock_server = MagicMock(spec=MCPServer) + mock_server.name = "filesystem" + + mock_server.shutdown = AsyncMock( + side_effect=RuntimeError("shutdown failed") + ) + + manager._servers["filesystem"] = mock_server + manager._started = True + + events = await collect_events(event_bus) + + # Shutdown failures are isolated and should not escape. + await manager.shutdown() + + failures = [ + event + for event in events + if isinstance( + event, + MCPServerShutdownFailedEvent, + ) + ] + + assert len(failures) == 1 + + assert failures[0].server_name == "filesystem" + assert failures[0].error == "shutdown failed" + + assert manager.servers == {} + assert manager.tools == [] + assert manager.started is False diff --git a/runtime/tests/integrations/test_mcp_server.py b/runtime/tests/integrations/test_mcp_server.py new file mode 100644 index 0000000..cfca023 --- /dev/null +++ b/runtime/tests/integrations/test_mcp_server.py @@ -0,0 +1,439 @@ +# runtime/tests/integrations/test_mcp_server.py + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from orion.integrations._mcp.config import MCPServerConfig +from orion.integrations._mcp.server import MCPServer + + +# ========================================================== +# Fixtures +# ========================================================== + + +@pytest.fixture +def stdio_config() -> MCPServerConfig: + return MCPServerConfig( + name="filesystem", + transport="stdio", + command="npx", + args=["-y", "test-mcp-server"], + env={"TEST_ENV": "true"}, + ) + + +@pytest.fixture +def http_config() -> MCPServerConfig: + return MCPServerConfig( + name="neo4j-mcp", + transport="http", + url="https://example.mcp.neo4j.io", + ) + + +@pytest.fixture +def server( + stdio_config: MCPServerConfig, +) -> MCPServer: + return MCPServer(stdio_config) + + +# ========================================================== +# Properties +# ========================================================== + + +def test_name( + server: MCPServer, +) -> None: + assert server.name == "filesystem" + + +def test_connected_initially_false( + server: MCPServer, +) -> None: + assert server.connected is False + + +# ========================================================== +# Startup - stdio +# ========================================================== + + +@pytest.mark.asyncio +async def test_startup_stdio( + stdio_config: MCPServerConfig, +) -> None: + + server = MCPServer(stdio_config) + + mock_read_stream = MagicMock() + mock_write_stream = MagicMock() + + mock_session = AsyncMock() + + with ( + patch( + "orion.integrations._mcp.server.stdio_client" + ) as mock_stdio_client, + patch.object( + server, + "_create_session", + new=AsyncMock(return_value=mock_session), + ) as mock_create_session, + ): + mock_stdio_client.return_value.__aenter__ = AsyncMock( + return_value=( + mock_read_stream, + mock_write_stream, + ) + ) + mock_stdio_client.return_value.__aexit__ = AsyncMock() + + await server.startup() + + assert server.connected is True + + mock_stdio_client.assert_called_once() + + mock_create_session.assert_awaited_once_with( + mock_read_stream, + mock_write_stream, + ) + + +@pytest.mark.asyncio +async def test_startup_stdio_requires_command() -> None: + + config = MCPServerConfig( + name="invalid", + transport="stdio", + command=None, + ) + + server = MCPServer(config) + + with pytest.raises( + ValueError, + match="uses stdio but no command was configured", + ): + await server.startup() + + assert server.connected is False + + +# ========================================================== +# Startup - HTTP +# ========================================================== + + +@pytest.mark.asyncio +async def test_startup_http( + http_config: MCPServerConfig, +) -> None: + + server = MCPServer(http_config) + + mock_read_stream = MagicMock() + mock_write_stream = MagicMock() + + mock_session = AsyncMock() + + with ( + patch( + "orion.integrations._mcp.server.streamable_http_client" + ) as mock_http_client, + patch.object( + server, + "_create_session", + new=AsyncMock(return_value=mock_session), + ) as mock_create_session, + ): + mock_http_client.return_value.__aenter__ = AsyncMock( + return_value=( + mock_read_stream, + mock_write_stream, + MagicMock(), + ) + ) + mock_http_client.return_value.__aexit__ = AsyncMock() + + await server.startup() + + assert server.connected is True + + mock_http_client.assert_called_once_with( + "https://example.mcp.neo4j.io" + ) + + mock_create_session.assert_awaited_once_with( + mock_read_stream, + mock_write_stream, + ) + + +@pytest.mark.asyncio +async def test_startup_http_requires_url() -> None: + + config = MCPServerConfig( + name="invalid", + transport="http", + url=None, + ) + + server = MCPServer(config) + + with pytest.raises( + ValueError, + match="uses HTTP but no URL was configured", + ): + await server.startup() + + assert server.connected is False + + +# ========================================================== +# Unsupported Transport +# ========================================================== + + +@pytest.mark.asyncio +async def test_startup_rejects_unsupported_transport() -> None: + + config = MCPServerConfig( + name="invalid", + transport="stdio", + command="echo", + ) + + server = MCPServer(config) + + # Bypass the Literal type at runtime to verify + # the defensive branch in startup(). + server.config = MagicMock( + name="invalid", + transport="websocket", + command=None, + args=[], + env={}, + url=None, + ) + + with pytest.raises( + ValueError, + match="Unsupported MCP transport 'websocket'", + ): + await server.startup() + + assert server.connected is False + + +# ========================================================== +# Startup Idempotency +# ========================================================== + + +@pytest.mark.asyncio +async def test_startup_does_nothing_when_already_connected( + server: MCPServer, +) -> None: + + server._connected = True + + with patch.object( + server, + "_startup_stdio", + new=AsyncMock(), + ) as startup_stdio: + + await server.startup() + + startup_stdio.assert_not_awaited() + assert server.connected is True + + +# ========================================================== +# Session Creation +# ========================================================== + + +@pytest.mark.asyncio +async def test_create_session( + server: MCPServer, +) -> None: + + mock_read_stream = MagicMock() + mock_write_stream = MagicMock() + + mock_session = AsyncMock() + + # The AsyncExitStack enters ClientSession as a context manager. + server._stack.enter_async_context = AsyncMock( + return_value=mock_session + ) + + result = await server._create_session( + mock_read_stream, + mock_write_stream, + ) + + assert result is mock_session + + mock_session.initialize.assert_awaited_once() + + server._stack.enter_async_context.assert_awaited_once() + + +# ========================================================== +# list_tools() +# ========================================================== + + +@pytest.mark.asyncio +async def test_list_tools_requires_connection( + server: MCPServer, +) -> None: + + with pytest.raises( + RuntimeError, + match="is not connected", + ): + await server.list_tools() + + +@pytest.mark.asyncio +async def test_list_tools( + server: MCPServer, +) -> None: + + mock_session = AsyncMock() + + expected = MagicMock() + + mock_session.list_tools.return_value = expected + + server._session = mock_session + server._connected = True + + result = await server.list_tools() + + assert result is expected + + mock_session.list_tools.assert_awaited_once() + + +# ========================================================== +# call_tool() +# ========================================================== + + +@pytest.mark.asyncio +async def test_call_tool_requires_connection( + server: MCPServer, +) -> None: + + with pytest.raises( + RuntimeError, + match="is not connected", + ): + await server.call_tool( + "test_tool", + {"value": "hello"}, + ) + + +@pytest.mark.asyncio +async def test_call_tool( + server: MCPServer, +) -> None: + + mock_session = AsyncMock() + + expected = MagicMock() + + mock_session.call_tool.return_value = expected + + server._session = mock_session + server._connected = True + + result = await server.call_tool( + "test_tool", + {"value": "hello"}, + ) + + assert result is expected + + mock_session.call_tool.assert_awaited_once_with( + "test_tool", + {"value": "hello"}, + ) + + +# ========================================================== +# Shutdown +# ========================================================== + + +@pytest.mark.asyncio +async def test_shutdown_when_not_connected( + server: MCPServer, +) -> None: + + server._connected = False + + server._stack.aclose = AsyncMock() + + await server.shutdown() + + server._stack.aclose.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shutdown( + server: MCPServer, +) -> None: + + server._connected = True + server._session = MagicMock() + + old_stack = server._stack + old_stack.aclose = AsyncMock() + + await server.shutdown() + + old_stack.aclose.assert_awaited_once() + + assert server.connected is False + assert server._session is None + + # A fresh stack is created so the server can be started again. + assert server._stack is not old_stack + + +@pytest.mark.asyncio +async def test_shutdown_cleans_state_when_stack_close_fails( + server: MCPServer, +) -> None: + + server._connected = True + server._session = MagicMock() + + old_stack = server._stack + + old_stack.aclose = AsyncMock( + side_effect=RuntimeError("close failed") + ) + + with pytest.raises( + RuntimeError, + match="close failed", + ): + await server.shutdown() + + assert server.connected is False + assert server._session is None + assert server._stack is not old_stack diff --git a/runtime/tests/llm/test_gemini.py b/runtime/tests/llm/test_gemini.py new file mode 100644 index 0000000..c47cc72 --- /dev/null +++ b/runtime/tests/llm/test_gemini.py @@ -0,0 +1,62 @@ +import pytest +from langchain_core.language_models import BaseChatModel +from langchain_google_genai import ChatGoogleGenerativeAI + +from orion.llm.config import LLMConfig +from orion.llm.providers.gemini import GeminiProvider + + +def test_create_gemini() -> None: + config = LLMConfig( + provider="gemini", + model="gemini-2.5-flash", + gemini_api_key="test-gemini-key", + ) + + provider = GeminiProvider(config) + + llm = provider.create() + + assert isinstance( + llm, + BaseChatModel, + ) + + assert isinstance( + llm, + ChatGoogleGenerativeAI, + ) + + assert llm.model == "gemini-2.5-flash" + + +def test_create_gemini_requires_api_key() -> None: + config = LLMConfig( + provider="gemini", + model="gemini-2.5-flash", + gemini_api_key="", + ) + + provider = GeminiProvider(config) + + with pytest.raises( + ValueError, + match="Gemini API key not set", + ): + provider.create() + + +def test_create_gemini_requires_model() -> None: + config = LLMConfig( + provider="gemini", + model="", + gemini_api_key="test-gemini-key", + ) + + provider = GeminiProvider(config) + + with pytest.raises( + ValueError, + match="Gemini model not set", + ): + provider.create() diff --git a/runtime/tests/llm/test_groq.py b/runtime/tests/llm/test_groq.py new file mode 100644 index 0000000..6720ebb --- /dev/null +++ b/runtime/tests/llm/test_groq.py @@ -0,0 +1,62 @@ +import pytest +from langchain_core.language_models import BaseChatModel +from langchain_groq import ChatGroq + +from orion.llm.config import LLMConfig +from orion.llm.providers.groq import GroqProvider + + +def test_create_groq() -> None: + config = LLMConfig( + provider="groq", + model="openai/gpt-oss-120b", + groq_api_key="test-groq-key", + ) + + provider = GroqProvider(config) + + llm = provider.create() + + assert isinstance( + llm, + BaseChatModel, + ) + + assert isinstance( + llm, + ChatGroq, + ) + + assert llm.model_name == "openai/gpt-oss-120b" + + +def test_create_groq_requires_api_key() -> None: + config = LLMConfig( + provider="groq", + model="openai/gpt-oss-120b", + groq_api_key="", + ) + + provider = GroqProvider(config) + + with pytest.raises( + ValueError, + match="Groq API key not set", + ): + provider.create() + + +def test_create_groq_requires_model() -> None: + config = LLMConfig( + provider="groq", + model="", + groq_api_key="test-groq-key", + ) + + provider = GroqProvider(config) + + with pytest.raises( + ValueError, + match="Groq model not set", + ): + provider.create() diff --git a/runtime/tests/memory/test_neo4j.py b/runtime/tests/memory/test_neo4j.py new file mode 100644 index 0000000..51d226a --- /dev/null +++ b/runtime/tests/memory/test_neo4j.py @@ -0,0 +1,777 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import pytest + +from orion.memory.models import Fact +from orion.memory.providers.graph.neo4j import ( + Neo4jKnowledgeGraph, +) + + +# ========================================================== +# Helpers +# ========================================================== + + +class AsyncResult: + def __init__( + self, + records: list[dict], + ) -> None: + self.records = records + + def __aiter__(self): + return self._iterate() + + async def _iterate(self): + for record in self.records: + yield record + + async def single(self): + if not self.records: + return None + + return self.records[0] + + +class FakeSession: + """ + Lightweight fake of a Neo4j async session. + + The real Neo4j driver supports both: + + session.run(query, parameters={...}) + + and: + + session.run(query, key=value) + + We preserve both forms here so the tests verify the actual + calls made by Neo4jKnowledgeGraph. + """ + + def __init__( + self, + results: list[AsyncResult] | None = None, + ) -> None: + self.results = results or [] + self.run_calls: list[tuple] = [] + self._result_index = 0 + + async def __aenter__(self): + return self + + async def __aexit__( + self, + exc_type, + exc, + traceback, + ): + return False + + async def run( + self, + query, + parameters=None, + **kwargs, + ): + self.run_calls.append( + ( + query, + parameters, + kwargs, + ) + ) + + if self._result_index < len(self.results): + result = self.results[self._result_index] + self._result_index += 1 + return result + + return AsyncResult([]) + + +class FakeDriver: + def __init__( + self, + session: FakeSession, + ) -> None: + self._session = session + + self.verify_connectivity = AsyncMock() + self.close = AsyncMock() + + def session(self): + return self._session + + +# ========================================================== +# Fixtures +# ========================================================== + + +@pytest.fixture +def graph() -> Neo4jKnowledgeGraph: + return Neo4jKnowledgeGraph( + uri="bolt://test:7687", + username="test-user", + password="test-password", + ) + + +@pytest.fixture +def session() -> FakeSession: + return FakeSession() + + +# ========================================================== +# Lifecycle +# ========================================================== + + +@pytest.mark.asyncio +async def test_startup_creates_driver( + graph, +) -> None: + driver = FakeDriver( + FakeSession(), + ) + + with patch( + "orion.memory.providers.graph.neo4j.AsyncGraphDatabase.driver", + return_value=driver, + ) as driver_factory: + await graph.startup() + + driver_factory.assert_called_once_with( + "bolt://test:7687", + auth=( + "test-user", + "test-password", + ), + ) + + driver.verify_connectivity.assert_awaited_once() + + assert graph.driver is driver + + +@pytest.mark.asyncio +async def test_shutdown_closes_driver( + graph, +) -> None: + driver = FakeDriver( + FakeSession(), + ) + + graph.driver = driver + + await graph.shutdown() + + driver.close.assert_awaited_once() + + assert graph.driver is None + + +@pytest.mark.asyncio +async def test_shutdown_without_driver( + graph, +) -> None: + await graph.shutdown() + + assert graph.driver is None + + +# ========================================================== +# Add Fact +# ========================================================== + + +@pytest.mark.asyncio +async def test_add_fact( + graph, + session, +) -> None: + graph.driver = FakeDriver(session) + + fact = Fact( + subject="Alice", + predicate="KNOWS", + object="Bob", + confidence=0.9, + ) + + await graph.add_fact(fact) + + assert len(session.run_calls) == 1 + + query, parameters, kwargs = session.run_calls[0] + + assert "MERGE (a:Entity {name: $subject})" in query + assert "MERGE (b:Entity {name: $object})" in query + assert "MERGE (a)-[r:KNOWS]->(b)" in query + assert "SET r.confidence = $confidence" in query + + assert parameters is None + + assert kwargs == { + "subject": "Alice", + "object": "Bob", + "confidence": 0.9, + } + + +@pytest.mark.asyncio +async def test_add_facts( + graph, +) -> None: + session = FakeSession() + + graph.driver = FakeDriver(session) + + facts = [ + Fact( + subject="Alice", + predicate="KNOWS", + object="Bob", + ), + Fact( + subject="Bob", + predicate="WORKS_AT", + object="Acme", + ), + ] + + await graph.add_facts(facts) + + assert len(session.run_calls) == 2 + + first_query, first_parameters, first_kwargs = ( + session.run_calls[0] + ) + + assert "MERGE (a)-[r:KNOWS]->(b)" in first_query + assert first_parameters is None + assert first_kwargs == { + "subject": "Alice", + "object": "Bob", + "confidence": 1.0, + } + + second_query, second_parameters, second_kwargs = ( + session.run_calls[1] + ) + + assert "MERGE (a)-[r:WORKS_AT]->(b)" in second_query + assert second_parameters is None + assert second_kwargs == { + "subject": "Bob", + "object": "Acme", + "confidence": 1.0, + } + + +@pytest.mark.asyncio +async def test_add_facts_empty( + graph, +) -> None: + session = FakeSession() + + graph.driver = FakeDriver(session) + + await graph.add_facts([]) + + assert session.run_calls == [] + + +# ========================================================== +# Remove Fact +# ========================================================== + + +@pytest.mark.asyncio +async def test_remove_fact( + graph, + session, +) -> None: + graph.driver = FakeDriver(session) + + fact = Fact( + subject="Alice", + predicate="KNOWS", + object="Bob", + ) + + await graph.remove_fact(fact) + + assert len(session.run_calls) == 1 + + query, parameters, kwargs = session.run_calls[0] + + assert "MATCH" in query + assert "(a:Entity {name: $subject})" in query + assert "[r:KNOWS]" in query + assert "(b:Entity {name: $object})" in query + assert "DELETE r" in query + + assert parameters is None + + assert kwargs == { + "subject": "Alice", + "object": "Bob", + } + + +@pytest.mark.asyncio +async def test_remove_facts( + graph, +) -> None: + session = FakeSession() + + graph.driver = FakeDriver(session) + + facts = [ + Fact( + subject="Alice", + predicate="KNOWS", + object="Bob", + ), + Fact( + subject="Bob", + predicate="WORKS_AT", + object="Acme", + ), + ] + + await graph.remove_facts(facts) + + assert len(session.run_calls) == 2 + + first_query, first_parameters, first_kwargs = ( + session.run_calls[0] + ) + + assert "[r:KNOWS]" in first_query + assert first_parameters is None + assert first_kwargs == { + "subject": "Alice", + "object": "Bob", + } + + second_query, second_parameters, second_kwargs = ( + session.run_calls[1] + ) + + assert "[r:WORKS_AT]" in second_query + assert second_parameters is None + assert second_kwargs == { + "subject": "Bob", + "object": "Acme", + } + + +@pytest.mark.asyncio +async def test_remove_facts_empty( + graph, +) -> None: + session = FakeSession() + + graph.driver = FakeDriver(session) + + await graph.remove_facts([]) + + assert session.run_calls == [] + + +# ========================================================== +# Search Facts +# ========================================================== + + +@pytest.mark.asyncio +async def test_search_facts( + graph, +) -> None: + result = AsyncResult( + [ + { + "subject": "Alice", + "predicate": "KNOWS", + "object": "Bob", + "confidence": 0.9, + }, + { + "subject": "Bob", + "predicate": "WORKS_AT", + "object": "Acme", + "confidence": 0.8, + }, + ] + ) + + session = FakeSession( + results=[result], + ) + + graph.driver = FakeDriver(session) + + facts = await graph.search_facts("Alice") + + assert facts == [ + Fact( + subject="Alice", + predicate="KNOWS", + object="Bob", + confidence=0.9, + ), + Fact( + subject="Bob", + predicate="WORKS_AT", + object="Acme", + confidence=0.8, + ), + ] + + query, parameters, kwargs = session.run_calls[0] + + assert "MATCH (a:Entity)-[r]->(b:Entity)" in query + assert "a.name = $query" in query + assert "b.name = $query" in query + + assert parameters == { + "query": "Alice", + } + + assert kwargs == {} + + +@pytest.mark.asyncio +async def test_search_facts_defaults_confidence( + graph, +) -> None: + result = AsyncResult( + [ + { + "subject": "Alice", + "predicate": "KNOWS", + "object": "Bob", + "confidence": None, + } + ] + ) + + session = FakeSession( + results=[result], + ) + + graph.driver = FakeDriver(session) + + facts = await graph.search_facts("Alice") + + assert len(facts) == 1 + assert facts[0].confidence == 1.0 + + +@pytest.mark.asyncio +async def test_search_facts_returns_empty( + graph, +) -> None: + session = FakeSession( + results=[ + AsyncResult([]), + ], + ) + + graph.driver = FakeDriver(session) + + facts = await graph.search_facts("Unknown") + + assert facts == [] + + +# ========================================================== +# Related Entities +# ========================================================== + + +@pytest.mark.asyncio +async def test_related_entities( + graph, +) -> None: + result = AsyncResult( + [ + { + "name": "Bob", + "label": "Person", + }, + { + "name": "Acme", + "label": "Company", + }, + ] + ) + + session = FakeSession( + results=[result], + ) + + graph.driver = FakeDriver(session) + + entities = await graph.related_entities( + "Alice", + depth=2, + ) + + assert len(entities) == 2 + + assert entities[0].name == "Bob" + assert entities[0].label == "Person" + + assert entities[1].name == "Acme" + assert entities[1].label == "Company" + + query, parameters, kwargs = session.run_calls[0] + + assert "[*1..2]" in query + + assert parameters is None + + assert kwargs == { + "entity": "Alice", + } + + +@pytest.mark.asyncio +async def test_related_entities_defaults_depth( + graph, +) -> None: + result = AsyncResult( + [ + { + "name": "Bob", + "label": "Person", + }, + ] + ) + + session = FakeSession( + results=[result], + ) + + graph.driver = FakeDriver(session) + + entities = await graph.related_entities("Alice") + + assert len(entities) == 1 + assert entities[0].name == "Bob" + + query, parameters, kwargs = session.run_calls[0] + + assert "[*1..1]" in query + + assert parameters is None + assert kwargs == { + "entity": "Alice", + } + + +@pytest.mark.asyncio +async def test_related_entities_depth_less_than_one( + graph, +) -> None: + session = FakeSession() + + graph.driver = FakeDriver(session) + + entities = await graph.related_entities( + "Alice", + depth=0, + ) + + assert entities == [] + assert session.run_calls == [] + + +@pytest.mark.asyncio +async def test_related_entities_defaults_missing_label( + graph, +) -> None: + result = AsyncResult( + [ + { + "name": "Bob", + "label": None, + }, + ] + ) + + session = FakeSession( + results=[result], + ) + + graph.driver = FakeDriver(session) + + entities = await graph.related_entities("Alice") + + assert len(entities) == 1 + assert entities[0].name == "Bob" + assert entities[0].label == "Entity" + + +# ========================================================== +# Schema +# ========================================================== + + +@pytest.mark.asyncio +async def test_get_schema( + graph, +) -> None: + labels_result = AsyncResult( + [ + {"label": "Entity"}, + {"label": "Person"}, + {"label": "Company"}, + ] + ) + + relationship_result = AsyncResult( + [ + {"relationshipType": "KNOWS"}, + {"relationshipType": "WORKS_AT"}, + ] + ) + + session = FakeSession( + results=[ + labels_result, + relationship_result, + ], + ) + + graph.driver = FakeDriver(session) + + schema = await graph.get_schema() + + assert schema.labels == [ + "Entity", + "Person", + "Company", + ] + + assert schema.relationship_types == [ + "KNOWS", + "WORKS_AT", + ] + + assert len(session.run_calls) == 2 + + labels_query, labels_parameters, labels_kwargs = ( + session.run_calls[0] + ) + + assert "CALL db.labels()" in labels_query + assert labels_parameters is None + assert labels_kwargs == {} + + relationship_query, relationship_parameters, relationship_kwargs = ( + session.run_calls[1] + ) + + assert "CALL db.relationshipTypes()" in relationship_query + assert relationship_parameters is None + assert relationship_kwargs == {} + + +@pytest.mark.asyncio +async def test_get_schema_empty( + graph, +) -> None: + session = FakeSession( + results=[ + AsyncResult([]), + AsyncResult([]), + ], + ) + + graph.driver = FakeDriver(session) + + schema = await graph.get_schema() + + assert schema.labels == [] + assert schema.relationship_types == [] + + assert len(session.run_calls) == 2 + + +# ========================================================== +# Maintenance +# ========================================================== + + +@pytest.mark.asyncio +async def test_clear( + graph, + session, +) -> None: + graph.driver = FakeDriver(session) + + await graph.clear() + + assert len(session.run_calls) == 1 + + query, parameters, kwargs = session.run_calls[0] + + assert query == "MATCH (n) DETACH DELETE n" + assert parameters is None + assert kwargs == {} + + +@pytest.mark.asyncio +async def test_count( + graph, +) -> None: + result = AsyncResult( + [ + { + "count": 5, + } + ] + ) + + session = FakeSession( + results=[result], + ) + + graph.driver = FakeDriver(session) + + count = await graph.count() + + assert count == 5 + + query, parameters, kwargs = session.run_calls[0] + + assert "MATCH ()-[r]->()" in query + assert "RETURN count(r) AS count" in query + + assert parameters is None + assert kwargs == {} + + +@pytest.mark.asyncio +async def test_count_empty_result( + graph, +) -> None: + session = FakeSession( + results=[ + AsyncResult([]), + ], + ) + + graph.driver = FakeDriver(session) + + count = await graph.count() + + assert count == 0 diff --git a/runtime/uv.lock b/runtime/uv.lock index 6048662..8c9d557 100644 --- a/runtime/uv.lock +++ b/runtime/uv.lock @@ -627,6 +627,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/60/02/be4a57b60c7149b55b9e3b3c13f609cd8eb5307c751f22bd8fb8d262e75b/filelock-3.29.7-py3-none-any.whl", hash = "sha256:987db6f789a3a2a59f55081801b2b3697cb97e2a736b5f1a9e99b559285fbc51", size = 46036, upload-time = "2026-07-08T05:46:57.53Z" }, ] +[[package]] +name = "filetype" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bb/29/745f7d30d47fe0f251d3ad3dc2978a23141917661998763bebb6da007eb1/filetype-1.2.0.tar.gz", hash = "sha256:66b56cd6474bf41d8c54660347d37afcc3f7d1970648de365c102ef77548aadb", size = 998020, upload-time = "2022-11-02T17:34:04.141Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/79/1b8fa1bb3568781e84c9200f951c735f3f157429f44be0495da55894d620/filetype-1.2.0-py2.py3-none-any.whl", hash = "sha256:7ce71b6880181241cf7ac8697a2f1eb6a8bd9b429f7ad6d27b8db9ba5f1c2d25", size = 19970, upload-time = "2022-11-02T17:34:01.425Z" }, +] + [[package]] name = "fsspec" version = "2026.6.0" @@ -636,6 +645,45 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/22/4222d7ddf3da30f363edaa98e329c2bce6c65497c9cb2810931c8b2c0fbc/fsspec-2026.6.0-py3-none-any.whl", hash = "sha256:02e0b71817df9b2169dc30a16832045764def1191b43dcff5bb85bdee212d2a1", size = 203949, upload-time = "2026-06-16T01:57:26.358Z" }, ] +[[package]] +name = "google-auth" +version = "2.56.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "pyasn1-modules" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/4c/fa42116a48bab3f7a143cf5042ecff7df9c8b73f8a376203cd534d1dc966/google_auth-2.56.3.tar.gz", hash = "sha256:40e229fc901f0a305b553050e5fce562d509bee0435be053abfa91582b51b90c", size = 367110, upload-time = "2026-08-06T06:24:01.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/b3/6117b2f24065cd7e2c4f140e9a193e215f089ca8ba314cf91eb9d0b7fe0a/google_auth-2.56.3-py3-none-any.whl", hash = "sha256:8ec438808f813ad034535000261eed1067475d229d05bbf4216e78c3f2362e53", size = 259116, upload-time = "2026-08-06T06:22:51.788Z" }, +] + +[package.optional-dependencies] +requests = [ + { name = "requests" }, +] + +[[package]] +name = "google-genai" +version = "2.19.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "google-auth", extra = ["requests"] }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "requests" }, + { name = "sniffio" }, + { name = "tenacity" }, + { name = "typing-extensions" }, + { name = "websockets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/37/1a/a834dfed90cf32dba900b533a1d14dcdefbda398bda5661d2a5a60fdc9fc/google_genai-2.19.0.tar.gz", hash = "sha256:d8f4126643793a7de230c396bcd142d21c948c8bb57507580e152549a7a41d9d", size = 659496, upload-time = "2026-08-19T23:05:43.276Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/e8/de0accd8cd004cf11252ca53fc5c3dda59bcf1856abfc7609842ceba7363/google_genai-2.19.0-py3-none-any.whl", hash = "sha256:36e0326dd886b52ef765be4c46042732b46b21f637abbe060e3db7c3de23974c", size = 1051056, upload-time = "2026-08-19T23:05:41.462Z" }, +] + [[package]] name = "groq" version = "0.37.1" @@ -981,9 +1029,10 @@ wheels = [ [[package]] name = "langchain-core" -version = "1.4.8" +version = "1.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "httpx" }, { name = "jsonpatch" }, { name = "langchain-protocol" }, { name = "langsmith" }, @@ -994,9 +1043,24 @@ dependencies = [ { name = "typing-extensions" }, { name = "uuid-utils" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/12/e3/bea6d0080acf183332f24dcd74c208aee5857cf8f783c3fb0bd86027d8fb/langchain_core-1.4.8.tar.gz", hash = "sha256:5bf1f8411077c904182ad8f975943d36adcbf579c4e017b3a118b719229ebf9a", size = 957974, upload-time = "2026-06-18T19:39:23.636Z" } +sdist = { url = "https://files.pythonhosted.org/packages/97/88/ebc98df187c525d729725ab39759337c56d5d2803423632376aa35bde899/langchain_core-1.6.0.tar.gz", hash = "sha256:dc72e36678ed26683ec0ad8829b44011fba461d10f9b31dbd9110215e5ec2333", size = 992493, upload-time = "2026-08-19T15:55:40.642Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/13/d6/bdf6f0481cc57ef300d6b1eb48cf1400c0409be715d6eb3cabadd1142a09/langchain_core-1.4.8-py3-none-any.whl", hash = "sha256:d84c28b05e3ba8d4271d0827aad5b592ccdaaf986e76768c23503f0a2045e8aa", size = 557416, upload-time = "2026-06-18T19:39:21.902Z" }, + { url = "https://files.pythonhosted.org/packages/84/4c/508a90b9d2e3bd7738fd93cb2ac2178ce734662c75e69b3b81c445f1b360/langchain_core-1.6.0-py3-none-any.whl", hash = "sha256:d8bb924cd413955d9d3192ccced140407427c3ff51e50ffb941222648da92cb2", size = 570006, upload-time = "2026-08-19T15:55:38.935Z" }, +] + +[[package]] +name = "langchain-google-genai" +version = "4.3.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filetype" }, + { name = "google-genai" }, + { name = "langchain-core" }, + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/41/e2/51eae10068e20917139f1d0b70bf2d5c15a8122d727457c1c50bed5a2d01/langchain_google_genai-4.3.5.tar.gz", hash = "sha256:bd33e6aa778f978c9bc632d59d5d6bcd1ed25fc07c2ed4aebc011a981c54f546", size = 289744, upload-time = "2026-08-20T16:31:44.51Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/98/93a37f508c8a5641389642883e70fba3e088a0b6a25f213445165a57c0cf/langchain_google_genai-4.3.5-py3-none-any.whl", hash = "sha256:24ed714b191c7bb85d2dba746aaab6eb8fd352fe782fdaefe50fde0d7b7601e9", size = 74207, upload-time = "2026-08-20T16:31:43.111Z" }, ] [[package]] @@ -1726,6 +1790,7 @@ dependencies = [ { name = "kokoro" }, { name = "langchain" }, { name = "langchain-core" }, + { name = "langchain-google-genai" }, { name = "langchain-groq" }, { name = "langchain-mcp-adapters" }, { name = "langgraph" }, @@ -1760,6 +1825,7 @@ requires-dist = [ { name = "kokoro", specifier = ">=0.9.4" }, { name = "langchain", specifier = ">=1.3.11" }, { name = "langchain-core", specifier = ">=1.4.8" }, + { name = "langchain-google-genai", specifier = ">=4.3.5" }, { name = "langchain-groq", specifier = ">=1.1.3" }, { name = "langchain-mcp-adapters", specifier = ">=0.3.0" }, { name = "langgraph", specifier = ">=1.2.8" }, @@ -2107,6 +2173,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" }, ] +[[package]] +name = "pyasn1" +version = "0.6.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a4/9a/23310166d960def5897e91fe20e5b724601b02a22e84ba1f94232c0b7f67/pyasn1-0.6.4.tar.gz", hash = "sha256:9c447d8431c947fe4c8febc4ed9e760bc29011a5b01e5c74b67025bd9fb8ce81", size = 151262, upload-time = "2026-07-09T01:12:33.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/3b/6163796d69c3977d1e4287bea4a6979161cbbdd170ebb430511e8e1999ce/pyasn1-0.6.4-py3-none-any.whl", hash = "sha256:deda9277cfd454080ec40b207fb6df82206a3a2688735233cdcd8d3d565f088b", size = 84410, upload-time = "2026-07-09T01:12:32.92Z" }, +] + +[[package]] +name = "pyasn1-modules" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyasn1" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892, upload-time = "2025-03-28T02:41:22.17Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259, upload-time = "2025-03-28T02:41:19.028Z" }, +] + [[package]] name = "pycparser" version = "3.0" From 2dc30b1050030c362458853aa7566ac20b03ec05 Mon Sep 17 00:00:00 2001 From: prasangeet Date: Tue, 25 Aug 2026 21:24:13 +0530 Subject: [PATCH 2/5] feat: major changes, updated mcp architecture and added multi-llm integration --- README.md | 188 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 188 insertions(+) diff --git a/README.md b/README.md index d380abd..c468e6b 100644 --- a/README.md +++ b/README.md @@ -448,6 +448,194 @@ This is the direction the project is heading: - Real token streaming and cancellable requests - Multi-modal perception and interaction +<<<<<<< HEAD +======= +## Repository Layout + +ORION is a monorepo with two independent applications that share a single +Git repository and communicate over an IPC protocol: + +- **`runtime/`** — the Python AI runtime (daemon). +- **`client/`** — the Rust (Ratatui) terminal client. + +```text +orion/ +├── runtime/ # Python AI runtime (daemon) +│ ├── src/orion/ # Installable application package +│ │ ├── __main__.py # Package entrypoint +│ │ ├── agent/ # Memory-aware agent graph and prompts +│ │ ├── bus/ # Event bus and subscription helpers +│ │ ├── cli/ # Typer CLI entrypoint +│ │ ├── core/ # Shared utilities such as the singleton metaclass +│ │ ├── events/ # Event models and registry +│ │ ├── integrations/ # External integrations such as MCP +│ │ ├── memory/ # Memory providers, planning, and persistence +│ │ ├── orchestrator/ # Runtime bootstrapping and pipeline entrypoint +│ │ ├── runtime/ # Runtime lifecycle and run loop +│ │ ├── services/ # Recording, STT, agent, TTS, playback, logging +│ │ ├── store/ # SQLite event persistence +│ │ └── transport/ # IPC protocol and bridge to the client +│ ├── tests/ # Behavior and smoke tests +│ ├── pyproject.toml # Runtime metadata and dependencies +│ └── uv.lock +├── client/ # Rust (Ratatui) terminal client +│ ├── Cargo.toml +│ └── src/main.rs +├── assets/ # Logo and visual assets +├── docker-compose.yml # Qdrant + Neo4j backends for memory +└── README.md +``` + +## How It Works + +### Event Bus + +The `EventBus` is the center of the runtime. Every event is written to the store first, then fanned out to subscribed handlers and global observers. + +### Orchestrator + +The orchestrator owns startup and shutdown: + +- builds the service list +- starts each service +- subscribes logging and UI observers +- emits a `PipelineStartEvent` +- tears everything down cleanly on exit + +### Services + +Each service is a small, focused unit: + +- `VoiceRecordingService` records microphone input and writes `data/audio/input.wav` +- `TranscriptGenerationService` sends audio to Groq Whisper and emits text +- `AgentService` turns the transcript into a response +- `TTSService` synthesizes speech into `data/audio/output.wav` +- `AudioPlaybackService` plays the response and closes the pipeline +- `TUIService` mirrors the live stream into the terminal UI +- `LoggingService` is available as a global observer hook + +## Requirements + +- Python 3.14+ +- Audio input device +- Audio output device +- `GROQ_API_KEY` configured in the environment + +## Installation + +The Python runtime lives in `runtime/`: + +```bash +cd runtime +uv sync +``` + +If you are not using `uv`, install the dependencies from `runtime/pyproject.toml` using your preferred Python tooling. + +## Configuration + +Create a `.env` file inside `runtime/` with your API key: + +```env +GROQ_API_KEY=your_key_here +``` + +Optional local artifacts: + +- `orion.db` - SQLite event store +- `data/audio/input.wav` - captured microphone input +- `data/audio/output.wav` - generated response audio + +## Run + +Start the memory backends (from the repo root), then run the runtime: + +```bash +docker compose up -d # Qdrant + Neo4j +cd runtime +uv run python -m orion # or: uv run orion +``` + +## Client (Rust) + +The terminal client lives in `client/` and is built with Cargo. Start the +runtime first (it opens the IPC socket the client connects to), then: + +```bash +cd client +cargo run +``` + +The client is a Ratatui TUI that talks to the runtime over `/tmp/orion.sock`. +It renders the conversation, a live event stream, and Copilot-style activity +logs, with tachyonfx animations. + +### Keybindings + +| Key | Action | +|-----|--------| +| `i` | enter insert mode (type a prompt) | +| `Enter` | send the typed prompt (insert mode) | +| `Esc` | back to normal mode | +| `v` | push-to-talk: press to start recording, press again to send | +| `s` | stop the assistant's speech | +| `j`/`k`, arrows, `PgUp`/`PgDn` | scroll the conversation | +| `q` | quit | + +### Voice + +Voice is handled entirely by the client; the runtime never touches audio +hardware: + +1. Press `v` to record from your microphone, `v` again to stop. +2. The client saves a WAV and sends its path to the runtime over IPC. +3. The runtime transcribes it (Groq Whisper) and runs the normal chat pipeline. +4. The response streams back and the client speaks it aloud. + +Text-to-speech uses your system speech engine via **speech-dispatcher**. Install +it to hear responses (otherwise TTS is silently skipped): + +```bash +sudo pacman -S speech-dispatcher # Arch / EndeavourOS +# Debian/Ubuntu: sudo apt-get install -y speech-dispatcher +``` + +Building the client also needs the ALSA development headers for microphone +capture (`cpal`): + +```bash +sudo pacman -S alsa-lib # Arch / EndeavourOS +# Debian/Ubuntu: sudo apt-get install -y libasound2-dev +``` + +## Tests + +```bash +cd runtime +uv run pytest +``` + +## Before Pushing a PR + +Run the project checks locally before you push a pull request: + +```bash +cd runtime +uv run pytest +uv run ruff check . +uv run ruff format --check . +``` + +If you are making code changes, it is worth running both commands again after your final edit so the PR starts clean. + +## Notes + +- `chat` is still a placeholder command. +- `doctor` is intentionally minimal right now. +- The current voice loop is continuous; it keeps starting new pipelines until the TUI exits. +- The codebase is already structured for more advanced agents, but desktop control and coding automation are still future work. + +>>>>>>> b709827 (docs: add client keybindings and voice usage (v/s keys, speech-dispatcher)) ## Vision ORION is being built toward a modular, observable intelligence system that can: From 27e4aff0e2b9d6ce018169ccd83b39b0d06836a9 Mon Sep 17 00:00:00 2001 From: prasangeet Date: Tue, 25 Aug 2026 21:22:23 +0530 Subject: [PATCH 3/5] feat: major changes, updated mcp architecture and added multi-llm integration --- README.md | 3 +++ client/src/ipc/events.rs | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/README.md b/README.md index c468e6b..d553c4e 100644 --- a/README.md +++ b/README.md @@ -448,6 +448,7 @@ This is the direction the project is heading: - Real token streaming and cancellable requests - Multi-modal perception and interaction +<<<<<<< HEAD <<<<<<< HEAD ======= ## Repository Layout @@ -636,6 +637,8 @@ If you are making code changes, it is worth running both commands again after yo - The codebase is already structured for more advanced agents, but desktop control and coding automation are still future work. >>>>>>> b709827 (docs: add client keybindings and voice usage (v/s keys, speech-dispatcher)) +======= +>>>>>>> 03ba8c4 (feat: major changes, updated mcp architecture and added multi-llm integration) ## Vision ORION is being built toward a modular, observable intelligence system that can: diff --git a/client/src/ipc/events.rs b/client/src/ipc/events.rs index 8023ccc..15beea5 100644 --- a/client/src/ipc/events.rs +++ b/client/src/ipc/events.rs @@ -61,7 +61,14 @@ pub enum RuntimeEvent { // ------------------------------------------------------------------ // Tools +<<<<<<< HEAD // ------------------------------------------------------------------ +======= +<<<<<<< HEAD +======= + // ------------------------------------------------------------------ +>>>>>>> b708882 (feat: major changes, updated mcp architecture and added multi-llm integration) +>>>>>>> 03ba8c4 (feat: major changes, updated mcp architecture and added multi-llm integration) ToolStarted { name: String, }, From 3807ab7fda8825b3949c46c27884b6a3eb4c0ee9 Mon Sep 17 00:00:00 2001 From: prasangeet Date: Tue, 25 Aug 2026 21:25:24 +0530 Subject: [PATCH 4/5] fix: fixed merge conflicts --- client/src/ipc/events.rs | 7 ------- 1 file changed, 7 deletions(-) diff --git a/client/src/ipc/events.rs b/client/src/ipc/events.rs index 15beea5..8023ccc 100644 --- a/client/src/ipc/events.rs +++ b/client/src/ipc/events.rs @@ -61,14 +61,7 @@ pub enum RuntimeEvent { // ------------------------------------------------------------------ // Tools -<<<<<<< HEAD // ------------------------------------------------------------------ -======= -<<<<<<< HEAD -======= - // ------------------------------------------------------------------ ->>>>>>> b708882 (feat: major changes, updated mcp architecture and added multi-llm integration) ->>>>>>> 03ba8c4 (feat: major changes, updated mcp architecture and added multi-llm integration) ToolStarted { name: String, }, From 955d1fd4588192480cb3cd672a7f4694ea831b43 Mon Sep 17 00:00:00 2001 From: prasangeet Date: Tue, 25 Aug 2026 21:27:32 +0530 Subject: [PATCH 5/5] fix: fixed merge conflicts --- README.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/README.md b/README.md index d553c4e..e0b4fd1 100644 --- a/README.md +++ b/README.md @@ -448,9 +448,6 @@ This is the direction the project is heading: - Real token streaming and cancellable requests - Multi-modal perception and interaction -<<<<<<< HEAD -<<<<<<< HEAD -======= ## Repository Layout ORION is a monorepo with two independent applications that share a single