From 6a2e059655182ca10d0f1a8ea042406f6b1026df Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Wed, 4 Mar 2026 22:04:19 +0800 Subject: [PATCH 01/21] refactor(sandbox): remove docker/aiosandbox backends, simplify SRT config - Remove docker and aiosandbox from available backends - Remove settings_path from SrtBackendConfig (now auto-generated in workspace) - Update SRT settings path to workspace/sandboxes/{session}-srt-settings.json - Update README examples to use srt backend and remove settingsPath --- bot/README.md | 5 +---- bot/vikingbot/config/schema.py | 1 - bot/vikingbot/sandbox/backends/srt.py | 4 ++-- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/bot/README.md b/bot/README.md index c8e88a5c..05fb0b6e 100644 --- a/bot/README.md +++ b/bot/README.md @@ -787,7 +787,7 @@ You only need to add sandbox configuration when you want to change these default ```json { "sandbox": { - "backend": "opensandbox", + "backend": "srt", "mode": "per-session" } } @@ -797,10 +797,8 @@ You only need to add sandbox configuration when you want to change these default | Backend | Description | |---------|-------------| | `direct` | (Default) Runs code directly on the host | -| `docker` | Uses Docker containers for isolation | | `opensandbox` | Uses OpenSandbox service | | `srt` | Uses Anthropic's SRT sandbox runtime | -| `aiosandbox` | Uses AIO Sandbox service | **Available Modes:** | Mode | Description | @@ -861,7 +859,6 @@ You only need to add sandbox configuration when you want to change these default "backend": "srt", "backends": { "srt": { - "settingsPath": "~/.vikingbot/srt-settings.json", "nodePath": "node", "network": { "allowedDomains": [], diff --git a/bot/vikingbot/config/schema.py b/bot/vikingbot/config/schema.py index 4bd1e10a..d5eddb8a 100644 --- a/bot/vikingbot/config/schema.py +++ b/bot/vikingbot/config/schema.py @@ -500,7 +500,6 @@ class DirectBackendConfig(BaseModel): class SrtBackendConfig(BaseModel): """SRT backend configuration.""" - settings_path: str = "~/.vikingbot/srt-settings.json" node_path: str = "node" network: SandboxNetworkConfig = Field(default_factory=SandboxNetworkConfig) filesystem: SandboxFilesystemConfig = Field(default_factory=SandboxFilesystemConfig) diff --git a/bot/vikingbot/sandbox/backends/srt.py b/bot/vikingbot/sandbox/backends/srt.py index eaff8c42..22e7a393 100644 --- a/bot/vikingbot/sandbox/backends/srt.py +++ b/bot/vikingbot/sandbox/backends/srt.py @@ -41,9 +41,9 @@ def _generate_settings(self) -> Path: """Generate SRT configuration file.""" srt_config = self._load_config() + # Place settings file in workspace/sandboxes/ directory settings_path = ( - Path.home() - / ".vikingbot" + self._workspace / "sandboxes" / f"{self.session_key.safe_name()}-srt-settings.json" ) From 9e7d877a42f9d12aa1dad23c40adba13c5199088 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Wed, 4 Mar 2026 22:10:27 +0800 Subject: [PATCH 02/21] refactor(sandbox): remove docker/aiosandbox backends, simplify SRT config - Remove docker and aiosandbox from available backends - Remove settings_path from SrtBackendConfig (now auto-generated in workspace) - Update SRT settings path to workspace/sandboxes/{session}-srt-settings.json - Update README examples to use srt backend and remove settingsPath --- bot/README_CN.md | 36 ------------- openviking_cli/cli/commands/chat.py | 82 ----------------------------- 2 files changed, 118 deletions(-) delete mode 100644 openviking_cli/cli/commands/chat.py diff --git a/bot/README_CN.md b/bot/README_CN.md index 48f87966..e1a101ed 100644 --- a/bot/README_CN.md +++ b/bot/README_CN.md @@ -475,10 +475,8 @@ vikingbot 支持沙箱执行以增强安全性。 | 后端 | 描述 | |---------|-------------| | `direct` | （默认）直接在主机上运行代码 | -| `docker` | 使用 Docker 容器进行隔离 | | `opensandbox` | 使用 OpenSandbox 服务 | | `srt` | 使用 Anthropic 的 SRT 沙箱运行时 | -| `aiosandbox` | 使用 AIO Sandbox 服务 | **可用模式：** | 模式 | 描述 | @@ -521,23 +519,6 @@ vikingbot 支持沙箱执行以增强安全性。 } ``` -**Docker 后端：** -```json -{ - "bot": { - "sandbox": { - "backend": "docker", - "backends": { - "docker": { - "image": "python:3.11-slim", - "networkMode": "bridge" - } - } - } - } -} -``` - **SRT 后端：** ```json { @@ -546,7 +527,6 @@ vikingbot 支持沙箱执行以增强安全性。 "backend": "srt", "backends": { "srt": { - "settingsPath": "~/.vikingbot/srt-settings.json", "nodePath": "node", "network": { "allowedDomains": [], @@ -569,22 +549,6 @@ vikingbot 支持沙箱执行以增强安全性。 } ``` -**AIO Sandbox 后端：** -```json -{ - "bot": { - "sandbox": { - "backend": "aiosandbox", - "backends": { - "aiosandbox": { - "baseUrl": "http://localhost:18794" - } - } - } - } -} -``` - **SRT 后端设置：** SRT 后端使用 `@anthropic-ai/sandbox-runtime`。 diff --git a/openviking_cli/cli/commands/chat.py b/openviking_cli/cli/commands/chat.py deleted file mode 100644 index 71c7b6ae..00000000 --- a/openviking_cli/cli/commands/chat.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 -"""Chat command - wrapper for vikingbot agent.""" - -import importlib.util -import shutil -import subprocess -import sys - -import typer - - -def _check_vikingbot() -> bool: - """Check if vikingbot is available.""" - return importlib.util.find_spec("vikingbot") is not None - - -def chat( - message: str = typer.Option(None, "--message", "-m", help="Message to send to the agent"), - session_id: str = typer.Option( - "cli__default__direct", "--session", "-s", help="Session ID" - ), - markdown: bool = typer.Option( - True, "--markdown/--no-markdown", help="Render assistant output as Markdown" - ), - logs: bool = typer.Option( - False, "--logs/--no-logs", help="Show vikingbot runtime logs during chat" - ), -): - """ - Chat with vikingbot agent. - - This is equivalent to `vikingbot chat`. - """ - if not _check_vikingbot(): - typer.echo( - typer.style( - "Error: vikingbot not found. Please install vikingbot first:", - fg="red", - ) - ) - typer.echo() - typer.echo(" Option 1: Install from local source (recommended for development)") - typer.echo(" cd bot") - typer.echo(" uv pip install -e \".[dev]\"") - typer.echo() - typer.echo(" Option 2: Install from PyPI (coming soon)") - typer.echo(" pip install vikingbot") - typer.echo() - raise typer.Exit(1) - - # Build the command arguments - args = [] - - if message: - args.extend(["--message", message]) - args.extend(["--session", session_id]) - if not markdown: - args.append("--no-markdown") - if logs: - args.append("--logs") - - # Check if vikingbot command exists - vikingbot_path = shutil.which("vikingbot") - - if vikingbot_path: - # Build the command: vikingbot chat [args...] - full_args = [vikingbot_path, "chat"] + args - else: - # Fallback: use python -m - full_args = [sys.executable, "-m", "vikingbot.cli.commands", "chat"] + args - - # Pass through all arguments to vikingbot agent - try: - subprocess.run(full_args, check=True) - except subprocess.CalledProcessError as e: - raise typer.Exit(e.returncode) - - -def register(app: typer.Typer) -> None: - """Register chat command.""" - app.command("chat")(chat) From e8aae41c69870d9d99719cde2b6a59819d26f7f2 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 11:38:25 +0800 Subject: [PATCH 03/21] fix: remove unused handle_chat_direct function and fix unused logs variable --- crates/ov_cli/src/main.rs | 73 +-------------------------------------- 1 file changed, 1 insertion(+), 72 deletions(-) diff --git a/crates/ov_cli/src/main.rs b/crates/ov_cli/src/main.rs index 72bdfeb1..4edffa90 100644 --- a/crates/ov_cli/src/main.rs +++ b/crates/ov_cli/src/main.rs @@ -477,14 +477,6 @@ enum ConfigCommands { #[tokio::main] async fn main() { - // Check for chat command first - handle it directly to bypass clap global args - // but we keep it in Cli enum so it shows up in help - let args: Vec = std::env::args().collect(); - if args.len() >= 2 && args[1] == "chat" { - handle_chat_direct(&args[2..]).await; - return; - } - let cli = Cli::parse(); let output_format = cli.output; @@ -583,7 +575,7 @@ async fn main() { Commands::Tui { uri } => { handle_tui(uri, ctx).await } - Commands::Chat { message, session, markdown, logs } => { + Commands::Chat { message, session, markdown, logs: _ } => { let cmd = commands::chat::ChatCommand { endpoint: std::env::var("VIKINGBOT_ENDPOINT").unwrap_or_else(|_| "http://localhost:18790/api/v1/openapi".to_string()), api_key: std::env::var("VIKINGBOT_API_KEY").ok(), @@ -1026,69 +1018,6 @@ async fn handle_health(ctx: CliContext) -> Result<()> { Ok(()) } -async fn handle_chat_direct(args: &[String]) { - use tokio::process::Command; - - // First check if vikingbot is available - let vikingbot_available = which::which("vikingbot").is_ok() || { - // Also check if we can import the module - let python = std::env::var("PYTHON").unwrap_or_else(|_| "python3".to_string()); - let check = Command::new(&python) - .args(["-c", "import vikingbot; print('ok')"]) - .output() - .await; - check.map(|o| o.status.success()).unwrap_or(false) - }; - - if !vikingbot_available { - eprintln!("Error: vikingbot not found. Please install vikingbot first:"); - eprintln!(); - eprintln!(" Option 1: Install from local source (recommended for development)"); - eprintln!(" cd bot"); - eprintln!(" uv pip install -e \".[dev]\""); - eprintln!(); - eprintln!(" Option 2: Install from PyPI (coming soon)"); - eprintln!(" pip install vikingbot"); - eprintln!(); - std::process::exit(1); - } - - // Try to find vikingbot executable first - let (cmd, mut vikingbot_args) = if let Ok(vikingbot) = which::which("vikingbot") { - (vikingbot, vec!["chat".to_string()]) - } else { - let python = std::env::var("PYTHON").unwrap_or_else(|_| "python3".to_string()); - ( - std::path::PathBuf::from(python), - vec!["-m".to_string(), "vikingbot.cli.commands".to_string(), "chat".to_string()], - ) - }; - - // Always add our default session first - vikingbot_args.push("--session".to_string()); - vikingbot_args.push("cli__chat__default".to_string()); - - // Now add all user args - if user provided --session it will override the default - vikingbot_args.extend(args.iter().cloned()); - - // Execute and pass through all signals - let status = Command::new(&cmd) - .args(&vikingbot_args) - .status() - .await; - - match status { - Ok(s) if !s.success() => { - std::process::exit(s.code().unwrap_or(1)); - } - Err(e) => { - eprintln!("Error: {}", e); - std::process::exit(1); - } - _ => {} - } -} - async fn handle_tui(uri: String, ctx: CliContext) -> Result<()> { let client = ctx.get_client(); tui::run_tui(client, &uri).await From 775ebd0f78d661a9744da5af815c834d5e2e7080 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 12:19:23 +0800 Subject: [PATCH 04/21] Fix UTF-8 issues in chat command --- crates/ov_cli/src/commands/chat.rs | 44 +++++++++++++++++++++++------- 1 file changed, 34 insertions(+), 10 deletions(-) diff --git a/crates/ov_cli/src/commands/chat.rs b/crates/ov_cli/src/commands/chat.rs index 23dcc940..0d6a4db7 100644 --- a/crates/ov_cli/src/commands/chat.rs +++ b/crates/ov_cli/src/commands/chat.rs @@ -1,15 +1,36 @@ //! Chat command for interacting with Vikingbot via OpenAPI -use std::io::Write; +use std::io::{BufRead, Write}; use std::time::Duration; +/// Safely truncate a string at a UTF-8 character boundary +fn truncate_utf8(s: &str, max_bytes: usize) -> &str { + if s.len() <= max_bytes { + return s; + } + + // Find the last valid UTF-8 character boundary before or at max_bytes + let mut boundary = max_bytes; + while boundary > 0 && !s.is_char_boundary(boundary) { + boundary -= 1; + } + + // If we couldn't find a boundary (unlikely), just return empty string + // Otherwise return up to the boundary + if boundary == 0 { + "" + } else { + &s[..boundary] + } +} + use clap::Parser; use reqwest::Client; use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; -const DEFAULT_ENDPOINT: &str = "http://localhost:18790/api/v1/openapi"; +const DEFAULT_ENDPOINT: &str = "http://localhost:1933/bot/v1"; /// Chat with Vikingbot via OpenAPI #[derive(Debug, Parser)] @@ -141,7 +162,7 @@ impl ChatCommand { "reasoning" => { let content = data.as_str().unwrap_or(""); if !self.no_format { - println!("\x1b[2mThink: {}...\x1b[0m", &content[..content.len().min(100)]); + println!("\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); } } "tool_call" => { @@ -154,7 +175,7 @@ impl ChatCommand { let content = data.as_str().unwrap_or(""); if !self.no_format { let truncated = if content.len() > 150 { - format!("{}...", &content[..150]) + format!("{}...", truncate_utf8(content, 150)) } else { content.to_string() }; @@ -194,11 +215,14 @@ impl ChatCommand { loop { // Read input print!("\x1b[1;32mYou:\x1b[0m "); - std::io::stdout().flush().map_err(|e| Error::Io(e))?; + std::io::stdout().flush()?; - let mut input = String::new(); - std::io::stdin().read_line(&mut input).map_err(|e| Error::Io(e))?; - let input = input.trim(); + // Read input as bytes first to handle invalid UTF-8 gracefully + let mut buf = Vec::new(); + let stdin = std::io::stdin(); + let mut reader = stdin.lock(); + reader.read_until(b'\n', &mut buf)?; + let input = String::from_utf8_lossy(&buf).trim().to_string(); if input.is_empty() { continue; @@ -248,7 +272,7 @@ impl ChatCommand { "reasoning" => { let content = data.as_str().unwrap_or(""); if content.len() > 100 { - println!("\x1b[2mThink: {}...\x1b[0m", &content[..100]); + println!("\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); } else { println!("\x1b[2mThink: {}\x1b[0m", content); } @@ -259,7 +283,7 @@ impl ChatCommand { "tool_result" => { let content = data.as_str().unwrap_or(""); let truncated = if content.len() > 150 { - format!("{}...", &content[..150]) + format!("{}...", truncate_utf8(content, 150)) } else { content.to_string() }; From cda6b55e9f522a119c30ae313fcccbe788863148 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 12:31:54 +0800 Subject: [PATCH 05/21] Add tab indentation to Think, Calling, and Result lines in CLI output --- crates/ov_cli/src/commands/chat.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/ov_cli/src/commands/chat.rs b/crates/ov_cli/src/commands/chat.rs index 0d6a4db7..782b540d 100644 --- a/crates/ov_cli/src/commands/chat.rs +++ b/crates/ov_cli/src/commands/chat.rs @@ -162,13 +162,13 @@ impl ChatCommand { "reasoning" => { let content = data.as_str().unwrap_or(""); if !self.no_format { - println!("\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); + println!("\t\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); } } "tool_call" => { let content = data.as_str().unwrap_or(""); if !self.no_format { - println!("\x1b[2m├─ Calling: {}\x1b[0m", content); + println!("\t\x1b[2m├─ Calling: {}\x1b[0m", content); } } "tool_result" => { @@ -179,7 +179,7 @@ impl ChatCommand { } else { content.to_string() }; - println!("\x1b[2m└─ Result: {}\x1b[0m", truncated); + println!("\t\x1b[2m└─ Result: {}\x1b[0m", truncated); } } _ => {} @@ -272,13 +272,13 @@ impl ChatCommand { "reasoning" => { let content = data.as_str().unwrap_or(""); if content.len() > 100 { - println!("\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); + println!("\t\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); } else { - println!("\x1b[2mThink: {}\x1b[0m", content); + println!("\t\x1b[2mThink: {}\x1b[0m", content); } } "tool_call" => { - println!("\x1b[2m├─ Calling: {}\x1b[0m", data.as_str().unwrap_or("")); + println!("\t\x1b[2m├─ Calling: {}\x1b[0m", data.as_str().unwrap_or("")); } "tool_result" => { let content = data.as_str().unwrap_or(""); @@ -287,7 +287,7 @@ impl ChatCommand { } else { content.to_string() }; - println!("\x1b[2m└─ Result: {}\x1b[0m", truncated); + println!("\t\x1b[2m└─ Result: {}\x1b[0m", truncated); } _ => {} } From 92a18760b44805f77ff21f3e5593cc4e6dea3be3 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 12:33:23 +0800 Subject: [PATCH 06/21] Add first release workflow --- bot/.github/workflows/release-first.yml | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 bot/.github/workflows/release-first.yml diff --git a/bot/.github/workflows/release-first.yml b/bot/.github/workflows/release-first.yml new file mode 100644 index 00000000..1e0819b3 --- /dev/null +++ b/bot/.github/workflows/release-first.yml @@ -0,0 +1,25 @@ +name: First Release to PyPI + +on: + workflow_dispatch: # 手动触发 + +jobs: + release: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install build dependencies + run: pip install build + + - name: Build package + run: python -m build + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + password: ${{ secrets.VIKINGBOT_PYPI_API_TOKEN }} From 8b4dee0f7a79021b9c0607711b3c0f08ef313d71 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 14:13:31 +0800 Subject: [PATCH 07/21] Update release workflow with correct working directory --- .github/workflows/release-vikingbot-first.yml | 29 +++ bot/vikingbot/channels/stdio.py | 179 ------------------ 2 files changed, 29 insertions(+), 179 deletions(-) create mode 100644 .github/workflows/release-vikingbot-first.yml delete mode 100644 bot/vikingbot/channels/stdio.py diff --git a/.github/workflows/release-vikingbot-first.yml b/.github/workflows/release-vikingbot-first.yml new file mode 100644 index 00000000..bbbd560a --- /dev/null +++ b/.github/workflows/release-vikingbot-first.yml @@ -0,0 +1,29 @@ +name: First Release to PyPI + +on: + workflow_dispatch: # 手动触发 + +jobs: + release: + runs-on: ubuntu-latest + defaults: + run: + working-directory: bot + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install build dependencies + run: pip install build + + - name: Build package + run: python -m build + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + password: ${{ secrets.VIKINGBOT_PYPI_API_TOKEN }} + packages-dir: bot/dist/ diff --git a/bot/vikingbot/channels/stdio.py b/bot/vikingbot/channels/stdio.py deleted file mode 100644 index 58a82443..00000000 --- a/bot/vikingbot/channels/stdio.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 -"""Stdio channel for vikingbot - communicates via stdin/stdout.""" - -import asyncio -import json -import sys -from pathlib import Path -from typing import Any - -from loguru import logger - -from vikingbot.bus.events import InboundMessage, OutboundMessage -from vikingbot.bus.queue import MessageBus -from vikingbot.channels.base import BaseChannel -from vikingbot.config.schema import SessionKey, BaseChannelConfig, ChannelType - - -class StdioChannelConfig(BaseChannelConfig): - """Configuration for StdioChannel.""" - - enabled: bool = True - type: Any = "stdio" - - def channel_id(self) -> str: - return "stdio" - - -class StdioChannel(BaseChannel): - """ - Stdio channel for vikingbot. - - This channel communicates via stdin/stdout using JSON messages: - - Reads JSON messages from stdin - - Publishes them to the MessageBus - - Subscribes to outbound messages and writes them to stdout - """ - - name: str = "stdio" - - def __init__( - self, config: BaseChannelConfig, bus: MessageBus, workspace_path: Path | None = None - ): - super().__init__(config, bus, workspace_path) - self._response_queue: asyncio.Queue[str] = asyncio.Queue() - - async def start(self) -> None: - """Start the stdio channel.""" - self._running = True - logger.info("Starting stdio channel") - - # Start reader and writer tasks - reader_task = asyncio.create_task(self._read_stdin()) - writer_task = asyncio.create_task(self._write_stdout()) - - # Send ready signal - await self._send_json({"type": "ready"}) - - try: - await asyncio.gather(reader_task, writer_task) - except asyncio.CancelledError: - self._running = False - reader_task.cancel() - writer_task.cancel() - await asyncio.gather(reader_task, writer_task, return_exceptions=True) - - async def stop(self) -> None: - """Stop the stdio channel.""" - self._running = False - logger.info("Stopping stdio channel") - - async def send(self, msg: OutboundMessage) -> None: - """Send a message via stdout.""" - if msg.is_normal_message: - await self._send_json({ - "type": "response", - "content": msg.content, - }) - else: - # For thinking events, just send the content as-is - await self._send_json({ - "type": "event", - "event_type": msg.event_type.value if hasattr(msg.event_type, "value") else str(msg.event_type), - "content": msg.content, - }) - - async def _send_json(self, data: dict[str, Any]) -> None: - """Send JSON data to stdout.""" - try: - line = json.dumps(data, ensure_ascii=False) - print(line, flush=True) - except Exception as e: - logger.exception(f"Failed to send JSON: {e}") - - async def _read_stdin(self) -> None: - """Read lines from stdin and publish to bus.""" - loop = asyncio.get_event_loop() - - while self._running: - try: - # Read a line from stdin - line = await loop.run_in_executor(None, sys.stdin.readline) - - if not line: - # EOF - self._running = False - break - - line = line.strip() - if not line: - continue - - # Parse the input - try: - request = json.loads(line) - except json.JSONDecodeError: - # Treat as simple text message - request = {"type": "message", "content": line} - - await self._handle_request(request) - - except Exception as e: - logger.exception(f"Error reading from stdin: {e}") - await self._send_json({ - "type": "error", - "message": str(e), - }) - - async def _write_stdout(self) -> None: - """Write responses from the queue to stdout.""" - while self._running: - try: - # Wait for a response with timeout - content = await asyncio.wait_for( - self._response_queue.get(), - timeout=0.5, - ) - await self._send_json({ - "type": "response", - "content": content, - }) - except asyncio.TimeoutError: - continue - except Exception as e: - logger.exception(f"Error writing to stdout: {e}") - - async def _handle_request(self, request: dict[str, Any]) -> None: - """Handle an incoming request.""" - request_type = request.get("type", "message") - - if request_type == "ping": - await self._send_json({"type": "pong"}) - - elif request_type == "message": - content = request.get("content", "") - chat_id = request.get("chat_id", "default") - sender_id = request.get("sender_id", "user") - - # Create and publish inbound message - msg = InboundMessage( - session_key=SessionKey( - type="stdio", - channel_id=self.channel_id, - chat_id=chat_id, - ), - sender_id=sender_id, - content=content, - ) - await self.bus.publish_inbound(msg) - - elif request_type == "quit": - await self._send_json({"type": "bye"}) - self._running = False - - else: - await self._send_json({ - "type": "error", - "message": f"Unknown request type: {request_type}", - }) From 6049dc0a9275c0d6ec4daaaffea32ea573a63fd3 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 14:58:01 +0800 Subject: [PATCH 08/21] =?UTF-8?q?=E4=BF=AE=E6=94=B9=20SessionKey=20?= =?UTF-8?q?=E6=9E=84=E5=BB=BA=E9=80=BB=E8=BE=91=EF=BC=9A=E7=BB=9F=E4=B8=80?= =?UTF-8?q?=E4=BD=BF=E7=94=A8=20type=3D"cli"=EF=BC=8Cchannel=5Fid=20?= =?UTF-8?q?=E9=BB=98=E8=AE=A4=20"default"=EF=BC=8Cchat=5Fid=20=E4=BD=9C?= =?UTF-8?q?=E4=B8=BA=20session=5Fid=20=E4=BD=BF=E7=94=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot/.github/workflows/release-first.yml | 25 -- bot/.github/workflows/release.yml | 33 ++ bot/scripts/clean_vikingbot.py | 155 -------- bot/scripts/clean_vikingbot.sh | 19 + bot/tests/conftest.py | 197 ++++++++++ bot/vikingbot/channels/chat.py | 11 +- bot/vikingbot/channels/discord.py | 5 +- bot/vikingbot/channels/manager.py | 1 + bot/vikingbot/channels/openapi.py | 7 +- bot/vikingbot/channels/openapi_models.py | 2 +- bot/vikingbot/channels/single_turn.py | 11 +- bot/vikingbot/channels/telegram.py | 10 +- bot/vikingbot/cli/commands.py | 25 +- bot/vikingbot/config/schema.py | 3 +- bot/vikingbot/openviking_mount/manager.py | 7 +- .../openviking_mount/session_integration.py | 9 +- bot/vikingbot/utils/helpers.py | 8 +- crates/ov_cli/Cargo.toml | 2 +- crates/ov_cli/src/commands/chat_v2.rs | 346 ++++++++++++++++++ crates/ov_cli/src/commands/mod.rs | 1 + crates/ov_cli/src/main.rs | 4 +- 21 files changed, 654 insertions(+), 227 deletions(-) delete mode 100644 bot/.github/workflows/release-first.yml create mode 100644 bot/.github/workflows/release.yml delete mode 100755 bot/scripts/clean_vikingbot.py create mode 100755 bot/scripts/clean_vikingbot.sh create mode 100644 bot/tests/conftest.py create mode 100644 crates/ov_cli/src/commands/chat_v2.rs diff --git a/bot/.github/workflows/release-first.yml b/bot/.github/workflows/release-first.yml deleted file mode 100644 index 1e0819b3..00000000 --- a/bot/.github/workflows/release-first.yml +++ /dev/null @@ -1,25 +0,0 @@ -name: First Release to PyPI - -on: - workflow_dispatch: # 手动触发 - -jobs: - release: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - uses: actions/setup-python@v5 - with: - python-version: '3.11' - - - name: Install build dependencies - run: pip install build - - - name: Build package - run: python -m build - - - name: Publish to PyPI - uses: pypa/gh-action-pypi-publish@release/v1 - with: - password: ${{ secrets.VIKINGBOT_PYPI_API_TOKEN }} diff --git a/bot/.github/workflows/release.yml b/bot/.github/workflows/release.yml new file mode 100644 index 00000000..1d10d7fe --- /dev/null +++ b/bot/.github/workflows/release.yml @@ -0,0 +1,33 @@ +name: Release to PyPI + +on: + push: + tags: + - 'v*' + +jobs: + release: + runs-on: ubuntu-latest + + permissions: + id-token: write + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install build dependencies + run: | + pip install build + + - name: Build package + run: | + python -m build + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/bot/scripts/clean_vikingbot.py b/bot/scripts/clean_vikingbot.py deleted file mode 100755 index 0689281e..00000000 --- a/bot/scripts/clean_vikingbot.py +++ /dev/null @@ -1,155 +0,0 @@ -#!/usr/bin/env python3 -""" -VikingBot一键清理脚本 - -清理以下内容： -- sessions/ - 会话文件 -- workspace/ - 工作空间文件 -- cron/ - 定时任务数据 -- 保留 config.json 配置文件 - -使用方法： - python clean_vikingbot.py # 交互式确认 - python clean_vikingbot.py --yes # 不确认直接删除 - python clean_vikingbot.py --dry-run # 预览删除内容，不实际删除 -""" - -import sys -import shutil -from pathlib import Path - - -def is_dry_run() -> bool: - """检查是否是预览模式""" - return "--dry-run" in sys.argv - - -def get_vikingbot_dir() -> Path: - """获取vikingbot数据目录""" - return Path.home() / ".vikingbot" - - -def confirm_action(message: str) -> bool: - """交互式确认""" - if "--yes" in sys.argv or "-y" in sys.argv: - return True - - try: - response = input(f"\n{message} (y/N): ").strip().lower() - return response in ["y", "yes"] - except (EOFError, KeyboardInterrupt): - print("\n\n⏭️ 跳过操作") - return False - - -def clean_directory(dir_path: Path, description: str) -> bool: - """清理指定目录""" - if not dir_path.exists(): - print(f" ℹ️ {description} 不存在，跳过") - return True - - if not confirm_action(f"确定要删除 {description} 吗？"): - print(f" ⏭️ 跳过删除 {description}") - return False - - if is_dry_run(): - print(f" [预览] 将删除 {description}") - return True - - try: - shutil.rmtree(dir_path) - print(f" ✅ 已删除 {description}") - return True - except Exception as e: - print(f" ❌ 删除 {description} 失败: {e}") - return False - - -def delete_file(file_path: Path, description: str) -> bool: - """删除指定文件""" - if not file_path.exists(): - return True - - if is_dry_run(): - print(f" [预览] 将删除 {description}") - return True - - try: - file_path.unlink() - print(f" ✅ 已删除 {description}") - return True - except Exception as e: - print(f" ❌ 删除 {description} 失败: {e}") - return False - - -def main(): - """主函数""" - print("=" * 60) - print("🧹 VikingBot 一键清理工具") - print("=" * 60) - - if is_dry_run(): - print("\n[预览模式] 不会实际删除任何文件") - - vikingbot_dir = get_vikingbot_dir() - - if not vikingbot_dir.exists(): - print(f"\n⚠️ VikingBot目录不存在: {vikingbot_dir}") - print(" 没有需要清理的内容") - return 0 - - print(f"\n📂 VikingBot目录: {vikingbot_dir}") - print("\n将清理以下内容：") - print(" 1. sessions/ - 会话文件") - print(" 2. workspace/ - 工作空间文件") - print(" 3. cron/ - 定时任务数据") - print(" 4. sandboxes/ - 沙箱数据") - print(" 5. bridge/ - Bridge数据") - print("\n⚠️ 注意: config.json 配置文件将保留") - - # 统计清理前的文件 - total_deleted = 0 - items_to_clean = [ - ("sessions", "sessions/ 会话目录"), - ("workspace", "workspace/ 工作空间"), - ("cron", "cron/ 定时任务数据"), - ("sandboxes", "sandboxes/ 沙箱目录"), - ("bridge", "bridge/ Bridge目录"), - ] - - print("\n" + "-" * 60) - for dir_name, description in items_to_clean: - dir_path = vikingbot_dir / dir_name - if clean_directory(dir_path, description): - total_deleted += 1 - - # 检查是否还有其他临时文件（配置备份文件默认不删除） - print("\n" + "-" * 60) - print("检查临时文件...") - - # 只显示配置备份文件，但不删除 - backup_files = list(vikingbot_dir.glob("config*.json")) - backup_files = [f for f in backup_files if f.name != "config.json"] - - if backup_files: - print(f"\n发现 {len(backup_files)} 个配置备份文件（不自动删除）:") - for f in backup_files: - print(f" - {f.name}") - print("\n💡 如需删除这些备份文件，请手动删除或修改脚本启用此功能") - - print("\n" + "=" * 60) - if is_dry_run(): - print("📋 [预览模式] 以上是将要删除的内容预览") - elif total_deleted > 0: - print(f"✅ 清理完成！共删除了 {total_deleted} 个目录") - else: - print("ℹ️ 没有需要清理的内容") - print("\n💡 下次运行 vikingbot 时，workspace 会自动重新初始化") - print("=" * 60) - - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/bot/scripts/clean_vikingbot.sh b/bot/scripts/clean_vikingbot.sh new file mode 100755 index 00000000..e4b1891f --- /dev/null +++ b/bot/scripts/clean_vikingbot.sh @@ -0,0 +1,19 @@ +#!/bin/bash +BOT_DIR="$HOME/.openviking/data/bot" + +echo "🧹 Cleaning VikingBot data directory..." +echo "📂 Cleaning contents of: $BOT_DIR" + +if [ -d "$BOT_DIR" ]; then + echo "🗑️ Deleting items:" + for item in "$BOT_DIR"/*; do + if [ -e "$item" ]; then + echo " - $(basename "$item")" + rm -rf "$item" + fi + done + echo "✅ Done!" +else + echo "⚠️ Directory does not exist: $BOT_DIR" +fi + diff --git a/bot/tests/conftest.py b/bot/tests/conftest.py new file mode 100644 index 00000000..c9cf65ca --- /dev/null +++ b/bot/tests/conftest.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Global test fixtures""" + +import asyncio +import shutil +from pathlib import Path +from typing import AsyncGenerator, Generator + +import pytest +import pytest_asyncio + +from openviking import AsyncOpenViking + +# Test data root directory +PROJECT_ROOT = Path(__file__).parent.parent +TEST_TMP_DIR = PROJECT_ROOT / "test_data" / "tmp" + + +@pytest.fixture(scope="session") +def event_loop(): + """Create session-level event loop""" + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +@pytest.fixture(scope="function") +def temp_dir() -> Generator[Path, None, None]: + """Create temp directory, auto-cleanup before and after test""" + shutil.rmtree(TEST_TMP_DIR, ignore_errors=True) + TEST_TMP_DIR.mkdir(parents=True, exist_ok=True) + yield TEST_TMP_DIR + + +@pytest.fixture(scope="function") +def test_data_dir(temp_dir: Path) -> Path: + """Create test data directory""" + data_dir = temp_dir / "data" + data_dir.mkdir(parents=True, exist_ok=True) + return data_dir + + +@pytest.fixture(scope="function") +def sample_text_file(temp_dir: Path) -> Path: + """Create sample text file""" + file_path = temp_dir / "sample.txt" + file_path.write_text("This is a sample text file for testing OpenViking.") + return file_path + + +@pytest.fixture(scope="function") +def sample_markdown_file(temp_dir: Path) -> Path: + """Create sample Markdown file""" + file_path = temp_dir / "sample.md" + file_path.write_text( + """# Sample Document + +## Introduction +This is a sample markdown document for testing OpenViking. + +## Features +- Feature 1: Resource management +- Feature 2: Semantic search +- Feature 3: Session management + +## Usage +Use this document to test various OpenViking functionalities. +""" + ) + return file_path + + +@pytest.fixture(scope="function") +def sample_skill_file(temp_dir: Path) -> Path: + """Create sample skill file in SKILL.md format""" + file_path = temp_dir / "sample_skill.md" + file_path.write_text( + """--- +name: sample-skill +description: A sample skill for testing OpenViking skill management +tags: + - test + - sample +--- + +# Sample Skill + +## Description +A sample skill for testing OpenViking skill management. + +## Usage +Use this skill when you need to test skill functionality. + +## Instructions +1. Step one: Initialize the skill +2. Step two: Execute the skill +3. Step three: Verify the result +""" + ) + return file_path + + +@pytest.fixture(scope="function") +def sample_directory(temp_dir: Path) -> Path: + """Create sample directory with multiple files""" + dir_path = temp_dir / "sample_dir" + dir_path.mkdir(parents=True, exist_ok=True) + + (dir_path / "file1.txt").write_text("Content of file 1 for testing.") + (dir_path / "file2.md").write_text("# File 2\nContent of file 2 for testing.") + + subdir = dir_path / "subdir" + subdir.mkdir() + (subdir / "file3.txt").write_text("Content of file 3 in subdir for testing.") + + return dir_path + + +@pytest.fixture(scope="function") +def sample_files(temp_dir: Path) -> list[Path]: + """Create multiple sample files for batch testing""" + files = [] + for i in range(3): + file_path = temp_dir / f"batch_file_{i}.md" + file_path.write_text( + f"""# Batch File {i} + +## Content +This is batch file number {i} for testing batch operations. + +## Keywords +- batch +- test +- file{i} +""" + ) + files.append(file_path) + return files + + +# ============ Client Fixtures ============ + + +@pytest_asyncio.fixture(scope="function") +async def client(test_data_dir: Path) -> AsyncGenerator[AsyncOpenViking, None]: + """Create initialized OpenViking client""" + await AsyncOpenViking.reset() + + client = AsyncOpenViking(path=str(test_data_dir)) + await client.initialize() + + yield client + + await client.close() + await AsyncOpenViking.reset() + + +@pytest_asyncio.fixture(scope="function") +async def uninitialized_client(test_data_dir: Path) -> AsyncGenerator[AsyncOpenViking, None]: + """Create uninitialized OpenViking client (for testing initialization flow)""" + await AsyncOpenViking.reset() + + client = AsyncOpenViking(path=str(test_data_dir)) + + yield client + + try: + await client.close() + except Exception: + pass + await AsyncOpenViking.reset() + + +@pytest_asyncio.fixture(scope="function") +async def client_with_resource_sync( + client: AsyncOpenViking, sample_markdown_file: Path +) -> AsyncGenerator[tuple[AsyncOpenViking, str], None]: + """Create client with resource (sync mode, wait for vectorization)""" + result = await client.add_resource( + path=str(sample_markdown_file), reason="Test resource", wait=True + ) + uri = result.get("root_uri", "") + + yield client, uri + + +@pytest_asyncio.fixture(scope="function") +async def client_with_resource( + client: AsyncOpenViking, sample_markdown_file: Path +) -> AsyncGenerator[tuple[AsyncOpenViking, str], None]: + """Create client with resource (async mode, no wait for vectorization)""" + result = await client.add_resource(path=str(sample_markdown_file), reason="Test resource") + uri = result.get("root_uri", "") + yield client, uri + diff --git a/bot/vikingbot/channels/chat.py b/bot/vikingbot/channels/chat.py index 7e8bd6b5..68408884 100644 --- a/bot/vikingbot/channels/chat.py +++ b/bot/vikingbot/channels/chat.py @@ -23,9 +23,10 @@ class ChatChannelConfig(BaseChannelConfig): enabled: bool = True type: Any = "cli" + _channel_id: str = "default" def channel_id(self) -> str: - return "chat" + return self._channel_id class ChatChannel(BaseChannel): @@ -44,7 +45,7 @@ def __init__( config: BaseChannelConfig, bus: MessageBus, workspace_path: Path | None = None, - session_id: str = "cli__chat__default", + session_id: str = "default", markdown: bool = True, logs: bool = False, ): @@ -144,7 +145,11 @@ def _exit_on_sigint(signum, frame): self._last_response = None msg = InboundMessage( - session_key=SessionKey.from_safe_name(self.session_id), + session_key=SessionKey( + type="cli", + channel_id=self.config.channel_id(), + chat_id=self.session_id, + ), sender_id="user", content=user_input, ) diff --git a/bot/vikingbot/channels/discord.py b/bot/vikingbot/channels/discord.py index b80ad495..74030063 100644 --- a/bot/vikingbot/channels/discord.py +++ b/bot/vikingbot/channels/discord.py @@ -204,11 +204,12 @@ async def _handle_message_create(self, payload: dict[str, Any]) -> None: content_parts = [content] if content else [] media_paths: list[str] = [] + from vikingbot.utils.helpers import get_media_path + if self.workspace_path: media_dir = self.workspace_path / "media" else: - # Fallback to ~/.vikingbot/media if workspace not available - media_dir = Path.home() / ".vikingbot" / "media" + media_dir = get_media_path() for attachment in payload.get("attachments") or []: url = attachment.get("url") diff --git a/bot/vikingbot/channels/manager.py b/bot/vikingbot/channels/manager.py index f61d6344..9da5892a 100644 --- a/bot/vikingbot/channels/manager.py +++ b/bot/vikingbot/channels/manager.py @@ -63,6 +63,7 @@ def add_channel_from_config( channel_config, self.bus, groq_api_key=additional_deps.get("groq_api_key"), + workspace_path=workspace_path, ) elif channel_config.type == ChannelType.FEISHU: diff --git a/bot/vikingbot/channels/openapi.py b/bot/vikingbot/channels/openapi.py index 2443ea32..1f583542 100644 --- a/bot/vikingbot/channels/openapi.py +++ b/bot/vikingbot/channels/openapi.py @@ -41,9 +41,10 @@ class OpenAPIChannelConfig(BaseChannelConfig): api_key: str = "" # If empty, no auth required allow_from: list[str] = [] max_concurrent_requests: int = 100 + _channel_id: str = "default" def channel_id(self) -> str: - return "openapi" + return self._channel_id class PendingResponse: @@ -293,7 +294,7 @@ async def _handle_chat(self, request: ChatRequest) -> ChatResponse: try: # Build session key session_key = SessionKey( - type="openapi", + type="cli", channel_id=self.config.channel_id(), chat_id=session_id, ) @@ -362,7 +363,7 @@ async def event_generator(): try: # Build session key and send message session_key = SessionKey( - type="openapi", + type="cli", channel_id=self.config.channel_id(), chat_id=session_id, ) diff --git a/bot/vikingbot/channels/openapi_models.py b/bot/vikingbot/channels/openapi_models.py index 0d4da5c8..c4d36bb1 100644 --- a/bot/vikingbot/channels/openapi_models.py +++ b/bot/vikingbot/channels/openapi_models.py @@ -38,7 +38,7 @@ class ChatRequest(BaseModel): """Request body for chat endpoint.""" message: str = Field(..., description="User message to send", min_length=1) - session_id: Optional[str] = Field(default=None, description="Session ID (optional, will create new if not provided)") + session_id: Optional[str] = Field(default="default", description="Session ID (optional, will create new if not provided)") user_id: Optional[str] = Field(default=None, description="User identifier (optional)") stream: bool = Field(default=False, description="Whether to stream the response") context: Optional[List[ChatMessage]] = Field(default=None, description="Additional context messages") diff --git a/bot/vikingbot/channels/single_turn.py b/bot/vikingbot/channels/single_turn.py index d4c44f09..69d28c28 100644 --- a/bot/vikingbot/channels/single_turn.py +++ b/bot/vikingbot/channels/single_turn.py @@ -20,9 +20,10 @@ class SingleTurnChannelConfig(BaseChannelConfig): enabled: bool = True type: Any = "cli" + _channel_id: str = "default" def channel_id(self) -> str: - return "chat" + return self._channel_id class SingleTurnChannel(BaseChannel): @@ -41,7 +42,7 @@ def __init__( bus: MessageBus, workspace_path: Path | None = None, message: str = "", - session_id: str = "cli__chat__default", + session_id: str = "default", markdown: bool = True, eval: bool = False, ): @@ -59,7 +60,11 @@ async def start(self) -> None: # Send the message msg = InboundMessage( - session_key=SessionKey.from_safe_name(self.session_id), + session_key=SessionKey( + type="cli", + channel_id=self.config.channel_id(), + chat_id=self.session_id, + ), sender_id="default", content=self.message, ) diff --git a/bot/vikingbot/channels/telegram.py b/bot/vikingbot/channels/telegram.py index 2e76a32d..9e979769 100644 --- a/bot/vikingbot/channels/telegram.py +++ b/bot/vikingbot/channels/telegram.py @@ -99,11 +99,13 @@ class TelegramChannel(BaseChannel): def __init__( self, - config: TelegramConfig, + config: TelegramChannelConfig, bus: MessageBus, groq_api_key: str = "", + workspace_path: Path | None = None, + **kwargs, ): - super().__init__(config, bus, **kwargs) + super().__init__(config, bus, workspace_path=workspace_path, **kwargs) self.config: TelegramChannelConfig = config self.groq_api_key = groq_api_key self._app: Application | None = None @@ -317,12 +319,12 @@ async def _on_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) # Save to workspace/media/ from pathlib import Path + from vikingbot.utils.helpers import get_media_path if self.workspace_path: media_dir = self.workspace_path / "media" else: - # Fallback to ~/.vikingbot/media if workspace not available - media_dir = Path.home() / ".vikingbot" / "media" + media_dir = get_media_path() media_dir.mkdir(parents=True, exist_ok=True) file_path = media_dir / f"{media_file.file_id[:16]}{ext}" diff --git a/bot/vikingbot/cli/commands.py b/bot/vikingbot/cli/commands.py index ebf95135..31d41bdd 100644 --- a/bot/vikingbot/cli/commands.py +++ b/bot/vikingbot/cli/commands.py @@ -276,7 +276,7 @@ async def run(): def prepare_agent_loop(config, bus, session_manager, cron, quiet: bool = False, eval: bool = False): - sandbox_parent_path = config.bot_data_path + sandbox_parent_path = config.workspace_path source_workspace_path = get_source_workspace_path() sandbox_manager = SandboxManager(config, sandbox_parent_path, source_workspace_path) if config.sandbox.backend == "direct": @@ -479,23 +479,13 @@ def _thinking_ctx(logs: bool): return console.status("[dim]vikingbot is thinking...[/dim]", spinner="dots") -def prepare_agent_channel(config, bus, mode: str, message: str | None, session_id: str, markdown: bool, logs: bool, eval: bool = False): +def prepare_agent_channel(config, bus, message: str | None, session_id: str, markdown: bool, logs: bool, eval: bool = False): """Prepare channel for agent command.""" from vikingbot.channels.chat import ChatChannel, ChatChannelConfig - from vikingbot.channels.stdio import StdioChannel, StdioChannelConfig from vikingbot.channels.single_turn import SingleTurnChannel, SingleTurnChannelConfig channels = ChannelManager(bus) - - if mode == "stdio": - channel_config = StdioChannelConfig() - channel = StdioChannel( - channel_config, - bus, - workspace_path=config.workspace_path, - ) - channels.add_channel(channel) - elif message is not None: + if message is not None: # Single message mode - use SingleTurnChannel for clean output channel_config = SingleTurnChannelConfig() channel = SingleTurnChannel( @@ -534,9 +524,6 @@ def chat( logs: bool = typer.Option( False, "--logs/--no-logs", help="Show vikingbot runtime logs during chat" ), - mode: str = typer.Option( - "direct", "--mode", help="Mode: direct (interactive), stdio (JSON IPC)" - ), eval: bool = typer.Option( False, "--eval", "-e", help="Run evaluation mode, output JSON results" ), @@ -559,9 +546,9 @@ def chat( is_single_turn = message is not None # Use unified default session ID if session_id is None: - session_id = "cli__chat__default" + session_id = "default" cron = prepare_cron(bus, quiet=is_single_turn) - channels = prepare_agent_channel(config, bus, mode, message, session_id, markdown, logs, eval) + channels = prepare_agent_channel(config, bus, message, session_id, markdown, logs, eval) agent_loop = prepare_agent_loop(config, bus, session_manager, cron, quiet=is_single_turn, eval=eval) async def run(): @@ -830,7 +817,7 @@ def cron_add( store_path = get_data_dir() / "cron" / "jobs.json" service = CronService(store_path) - session_key = SessionKey.from_safe_name() + session_key = SessionKey(type="cli", channel_id="default", chat_id="default") job = service.add_job( name=name, diff --git a/bot/vikingbot/config/schema.py b/bot/vikingbot/config/schema.py index d5eddb8a..ef07f802 100644 --- a/bot/vikingbot/config/schema.py +++ b/bot/vikingbot/config/schema.py @@ -250,9 +250,10 @@ class OpenAPIChannelConfig(BaseChannelConfig): api_key: str = "" # If empty, no auth required allow_from: list[str] = Field(default_factory=list) max_concurrent_requests: int = 100 + _channel_id: str = "default" def channel_id(self) -> str: - return "openapi" + return self._channel_id class ChannelsConfig(BaseModel): diff --git a/bot/vikingbot/openviking_mount/manager.py b/bot/vikingbot/openviking_mount/manager.py index f5375422..0fa4c1f6 100644 --- a/bot/vikingbot/openviking_mount/manager.py +++ b/bot/vikingbot/openviking_mount/manager.py @@ -13,6 +13,7 @@ from loguru import logger +from vikingbot.utils.helpers import get_mounts_path, get_bot_data_path from .mount import OpenVikingMount, MountConfig, MountScope @@ -41,8 +42,8 @@ def __init__(self, base_mount_dir: Optional[Path] = None): base_mount_dir: 基础挂载目录，所有挂载点将在此目录下创建 """ if base_mount_dir is None: - # 默认在用户目录下创建 - base_mount_dir = Path.home() / ".vikingbot" / "mounts" + # 默认从配置路径获取 + base_mount_dir = get_mounts_path() self.base_mount_dir = base_mount_dir self._mounts: Dict[str, MountPoint] = {} @@ -218,7 +219,7 @@ def create_resources_mount( """ if openviking_data_path is None: # 默认使用vikingbot的openviking数据目录 - openviking_data_path = Path.home() / ".vikingbot" / "openviking_data" + openviking_data_path = get_bot_data_path() / "ov_data" return self.create_mount( mount_id=mount_id, diff --git a/bot/vikingbot/openviking_mount/session_integration.py b/bot/vikingbot/openviking_mount/session_integration.py index 2ccd68a9..b16b3f4d 100644 --- a/bot/vikingbot/openviking_mount/session_integration.py +++ b/bot/vikingbot/openviking_mount/session_integration.py @@ -1,7 +1,7 @@ """ OpenViking FUSE 会话集成 -提供与会话管理器的集成，自动在 .vikingbot/workspace/{session}/ 挂载 OpenViking +提供与会话管理器的集成，自动在配置的 workspace/{session}/ 挂载 OpenViking 每个session直接在自己的workspace下管理内容 """ @@ -15,6 +15,7 @@ from loguru import logger +from vikingbot.utils.helpers import get_workspace_path # 相对导入同一包内的模块 from .mount import OpenVikingMount, MountConfig, MountScope from .viking_fuse import mount_fuse, FUSEMountManager, FUSE_AVAILABLE @@ -35,7 +36,7 @@ def __init__(self, base_workspace: Optional[Path] = None): base_workspace: 基础工作区路径 """ if base_workspace is None: - base_workspace = Path.home() / ".vikingbot" / "workspace" + base_workspace = get_workspace_path() self.base_workspace = base_workspace self.base_workspace.mkdir(parents=True, exist_ok=True) @@ -58,7 +59,7 @@ def get_session_workspace(self, session_key: str) -> Path: session_key: 会话键 Returns: - 工作区路径: .vikingbot/workspace/{session}/ + 工作区路径: {workspace}/.vikingbot/workspace/{session}/ """ safe_session_key = session_key.replace(":", "__") return self.base_workspace / safe_session_key @@ -71,7 +72,7 @@ def get_session_ov_data_path(self, session_key: str) -> Path: session_key: 会话键 Returns: - 数据存储路径: .vikingbot/workspace/{session}/.ov_data/ + 数据存储路径: {workspace}/{session}/.ov_data/ """ return self.get_session_workspace(session_key) / ".ov_data" diff --git a/bot/vikingbot/utils/helpers.py b/bot/vikingbot/utils/helpers.py index cf1d761c..081ac0cf 100644 --- a/bot/vikingbot/utils/helpers.py +++ b/bot/vikingbot/utils/helpers.py @@ -2,10 +2,13 @@ from pathlib import Path from datetime import datetime +from loguru import logger def ensure_dir(path: Path) -> Path: """Ensure a directory exists, creating it if necessary.""" + if not path.exists(): + logger.info(f"Creating directory: {path}") path.mkdir(parents=True, exist_ok=True) return path @@ -17,7 +20,10 @@ def ensure_dir(path: Path) -> Path: def set_bot_data_path(path: Path) -> None: """Set the global bot data path.""" global _bot_data_path - _bot_data_path = path + expanded_path = path.expanduser() + if not expanded_path.exists(): + logger.info(f"Storage workspace directory does not exist, will be created: {expanded_path}") + _bot_data_path = expanded_path def get_bot_data_path() -> Path: diff --git a/crates/ov_cli/Cargo.toml b/crates/ov_cli/Cargo.toml index 8a93a1d8..08b74f03 100644 --- a/crates/ov_cli/Cargo.toml +++ b/crates/ov_cli/Cargo.toml @@ -29,4 +29,4 @@ zip = "2.2" tempfile = "3.12" url = "2.5" walkdir = "2.5" -which = "6.0" +rustyline = "14.0" diff --git a/crates/ov_cli/src/commands/chat_v2.rs b/crates/ov_cli/src/commands/chat_v2.rs new file mode 100644 index 00000000..afafcb6e --- /dev/null +++ b/crates/ov_cli/src/commands/chat_v2.rs @@ -0,0 +1,346 @@ + +//! Chat command for interacting with Vikingbot via OpenAPI (v2 with rustyline) + +use std::time::Duration; + +use clap::Parser; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use rustyline::error::ReadlineError; +use rustyline::{Editor, history::FileHistory}; + +use crate::error::{Error, Result}; + +const DEFAULT_ENDPOINT: &str = "http://localhost:1933/bot/v1"; + +/// Safely truncate a string at a UTF-8 character boundary +fn truncate_utf8(s: &str, max_bytes: usize) -> &str { + if s.len() <= max_bytes { + return s; + } + + let mut boundary = max_bytes; + while boundary > 0 && !s.is_char_boundary(boundary) { + boundary -= 1; + } + + if boundary == 0 { + "" + } else { + &s[..boundary] + } +} + +#[derive(Debug, Parser)] +pub struct ChatCommand { + #[arg(short, long, default_value = DEFAULT_ENDPOINT)] + pub endpoint: String, + + #[arg(short, long, env = "VIKINGBOT_API_KEY")] + pub api_key: Option, + + #[arg(short, long)] + pub session: Option, + + #[arg(short, long, default_value = "cli_user")] + pub user: String, + + #[arg(short = 'M', long)] + pub message: Option, + + #[arg(long)] + pub stream: bool, + + #[arg(long)] + pub no_format: bool, +} + +#[derive(Debug, Serialize, Deserialize)] +struct ChatMessage { + role: String, + content: String, +} + +#[derive(Debug, Serialize)] +struct ChatRequest { + message: String, + #[serde(skip_serializing_if = "Option::is_none")] + session_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + user_id: Option, + stream: bool, + #[serde(skip_serializing_if = "Option::is_none")] + context: Option>, +} + +#[derive(Debug, Deserialize)] +struct ChatResponse { + session_id: String, + message: String, + #[serde(default)] + events: Option>, +} + +#[derive(Debug, Deserialize)] +struct StreamEvent { + event: String, + data: serde_json::Value, +} + +impl ChatCommand { + pub async fn execute(&self) -> Result<()> { + let client = Client::builder() + .timeout(Duration::from_secs(300)) + .build() + .map_err(|e| Error::Network(format!("Failed to create HTTP client: {}", e)))?; + + if let Some(message) = &self.message { + self.send_message(&client, message).await + } else { + self.run_interactive(&client).await + } + } + + async fn send_message(&self, client: &Client, message: &str) -> Result<()> { + let url = format!("{}/chat", self.endpoint); + + let request = ChatRequest { + message: message.to_string(), + session_id: self.session.clone(), + user_id: Some(self.user.clone()), + stream: false, + context: None, + }; + + let mut req_builder = client.post(&url).json(&request); + + if let Some(api_key) = &self.api_key { + req_builder = req_builder.header("X-API-Key", api_key); + } + + let response = req_builder + .send() + .await + .map_err(|e| Error::Network(format!("Failed to send request: {}", e)))?; + + if !response.status().is_success() { + let status = response.status(); + let text = response.text().await.unwrap_or_default(); + return Err(Error::Api(format!("Request failed ({}): {}", status, text))); + } + + let chat_response: ChatResponse = response + .json() + .await + .map_err(|e| Error::Parse(format!("Failed to parse response: {}", e)))?; + + if let Some(events) = &chat_response.events { + for event in events { + if let (Some(etype), Some(data)) = ( + event.get("type").and_then(|v| v.as_str()), + event.get("data"), + ) { + match etype { + "reasoning" => { + let content = data.as_str().unwrap_or(""); + if !self.no_format { + println!("\t\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); + } + } + "tool_call" => { + let content = data.as_str().unwrap_or(""); + if !self.no_format { + println!("\t\x1b[2m├─ Calling: {}\x1b[0m", content); + } + } + "tool_result" => { + let content = data.as_str().unwrap_or(""); + if !self.no_format { + let truncated = if content.len() > 150 { + format!("{}...", truncate_utf8(content, 150)) + } else { + content.to_string() + }; + println!("\t\x1b[2m└─ Result: {}\x1b[0m", truncated); + } + } + _ => {} + } + } + } + } + + if !self.no_format { + println!("\n\x1b[1;31mBot:\x1b[0m"); + println!("{}", chat_response.message); + println!(); + } else { + println!("{}", chat_response.message); + } + + Ok(()) + } + + async fn run_interactive(&self, client: &Client) -> Result<()> { + println!("Vikingbot Chat - Interactive Mode (v2)"); + println!("Endpoint: {}", self.endpoint); + if let Some(session) = &self.session { + println!("Session: {}", session); + } + println!("Type 'exit', 'quit', or press Ctrl+C to exit"); + println!("----------------------------------------\n"); + + let mut session_id = self.session.clone(); + + let mut rl: Editor<(), FileHistory> = Editor::new() + .map_err(|e| Error::Client(format!("Failed to create line editor: {}", e)))?; + + let _ = rl.load_history("ov_chat_history.txt"); + + loop { + let readline = rl.readline("\x1b[1;32mYou:\x1b[0m "); + + let input = match readline { + Ok(line) => { + let trimmed = line.trim().to_string(); + if !trimmed.is_empty() { + let _ = rl.add_history_entry(line); + let _ = rl.save_history("ov_chat_history.txt"); + } + trimmed + } + Err(ReadlineError::Interrupted) => { + println!("\nGoodbye!"); + break; + } + Err(ReadlineError::Eof) => { + println!("\nGoodbye!"); + break; + } + Err(err) => { + eprintln!("\x1b[1;31mError reading input: {}\x1b[0m", err); + continue; + } + }; + + if input.is_empty() { + continue; + } + + if input.eq_ignore_ascii_case("exit") || input.eq_ignore_ascii_case("quit") { + println!("\nGoodbye!"); + break; + } + + let url = format!("{}/chat", self.endpoint); + + let request = ChatRequest { + message: input.to_string(), + session_id: session_id.clone(), + user_id: Some(self.user.clone()), + stream: false, + context: None, + }; + + let mut req_builder = client.post(&url).json(&request); + + if let Some(api_key) = &self.api_key { + req_builder = req_builder.header("X-API-Key", api_key); + } + + match req_builder.send().await { + Ok(response) => { + if response.status().is_success() { + match response.json::().await { + Ok(chat_response) => { + if session_id.is_none() { + session_id = Some(chat_response.session_id.clone()); + } + + if let Some(events) = chat_response.events { + for event in events { + if let (Some(etype), Some(data)) = ( + event.get("type").and_then(|v| v.as_str()), + event.get("data"), + ) { + match etype { + "reasoning" => { + let content = data.as_str().unwrap_or(""); + if content.len() > 100 { + println!("\t\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); + } else { + println!("\t\x1b[2mThink: {}\x1b[0m", content); + } + } + "tool_call" => { + println!("\t\x1b[2m├─ Calling: {}\x1b[0m", data.as_str().unwrap_or("")); + } + "tool_result" => { + let content = data.as_str().unwrap_or(""); + let truncated = if content.len() > 150 { + format!("{}...", truncate_utf8(content, 150)) + } else { + content.to_string() + }; + println!("\t\x1b[2m└─ Result: {}\x1b[0m", truncated); + } + _ => {} + } + } + } + } + + println!("\n\x1b[1;31mBot:\x1b[0m"); + println!("{}", chat_response.message); + println!(); + } + Err(e) => { + eprintln!("\x1b[1;31mError parsing response: {}\x1b[0m", e); + } + } + } else { + let status = response.status(); + let text = response.text().await.unwrap_or_default(); + eprintln!("\x1b[1;31mRequest failed ({}): {}\x1b[0m", status, text); + } + } + Err(e) => { + eprintln!("\x1b[1;31mFailed to send request: {}\x1b[0m", e); + } + } + } + + println!("\nGoodbye!"); + Ok(()) + } +} + +impl ChatCommand { + pub async fn run(&self) -> Result<()> { + self.execute().await + } +} + +impl ChatCommand { + #[allow(clippy::too_many_arguments)] + pub fn new( + endpoint: String, + api_key: Option, + session: Option, + user: String, + message: Option, + stream: bool, + no_format: bool, + ) -> Self { + Self { + endpoint, + api_key, + session, + user, + message, + stream, + no_format, + } + } +} + diff --git a/crates/ov_cli/src/commands/mod.rs b/crates/ov_cli/src/commands/mod.rs index 1e8d1bcd..270766a6 100644 --- a/crates/ov_cli/src/commands/mod.rs +++ b/crates/ov_cli/src/commands/mod.rs @@ -1,5 +1,6 @@ pub mod admin; pub mod chat; +pub mod chat_v2; pub mod content; pub mod search; pub mod filesystem; diff --git a/crates/ov_cli/src/main.rs b/crates/ov_cli/src/main.rs index 4edffa90..2c276269 100644 --- a/crates/ov_cli/src/main.rs +++ b/crates/ov_cli/src/main.rs @@ -576,8 +576,8 @@ async fn main() { handle_tui(uri, ctx).await } Commands::Chat { message, session, markdown, logs: _ } => { - let cmd = commands::chat::ChatCommand { - endpoint: std::env::var("VIKINGBOT_ENDPOINT").unwrap_or_else(|_| "http://localhost:18790/api/v1/openapi".to_string()), + let cmd = commands::chat_v2::ChatCommand { + endpoint: std::env::var("VIKINGBOT_ENDPOINT").unwrap_or_else(|_| "http://localhost:1933/bot/v1".to_string()), api_key: std::env::var("VIKINGBOT_API_KEY").ok(), session: Some(session), user: "cli_user".to_string(), From 0557391571a7b96834dabea7acd99d039b437e30 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 16:53:17 +0800 Subject: [PATCH 09/21] Implement machine unique ID as default session ID for ov chat --- bot/vikingbot/channels/openapi.py | 2 +- crates/ov_cli/src/config.rs | 55 +++++++++++++++++++++++++++++++ crates/ov_cli/src/main.rs | 11 ++++--- 3 files changed, 62 insertions(+), 6 deletions(-) diff --git a/bot/vikingbot/channels/openapi.py b/bot/vikingbot/channels/openapi.py index 1f583542..20afb1d1 100644 --- a/bot/vikingbot/channels/openapi.py +++ b/bot/vikingbot/channels/openapi.py @@ -437,7 +437,7 @@ def get_openapi_router(bus: MessageBus, config: Config) -> APIRouter: # Register channel's send method as subscriber for outbound messages bus.subscribe_outbound( - f"openapi__{openapi_config.channel_id()}", + f"cli__{openapi_config.channel_id()}", channel.send, ) diff --git a/crates/ov_cli/src/config.rs b/crates/ov_cli/src/config.rs index 57a81789..f7ba3a1c 100644 --- a/crates/ov_cli/src/config.rs +++ b/crates/ov_cli/src/config.rs @@ -98,3 +98,58 @@ pub fn default_config_path() -> Result { .ok_or_else(|| Error::Config("Could not determine home directory".to_string()))?; Ok(home.join(".openviking").join("ovcli.conf")) } + +/// Get or create a unique machine ID. +/// +/// This function will: +/// 1. Try to get the hostname from environment variables +/// 2. If that fails, generate a random ID and store it in the config directory +pub fn get_or_create_machine_id() -> Result { + let home = dirs::home_dir() + .ok_or_else(|| Error::Config("Could not determine home directory".to_string()))?; + let machine_id_path = home.join(".openviking").join("machine_id"); + + // Try to read existing machine ID + if machine_id_path.exists() { + if let Ok(content) = std::fs::read_to_string(&machine_id_path) { + let trimmed = content.trim(); + if !trimmed.is_empty() { + return Ok(trimmed.to_string()); + } + } + } + + // Try to get hostname from environment variables + if let Ok(hostname) = std::env::var("HOSTNAME").or_else(|_| std::env::var("COMPUTERNAME")) { + let trimmed = hostname.trim(); + if !trimmed.is_empty() { + // Save it for future use + if let Some(parent) = machine_id_path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let _ = std::fs::write(&machine_id_path, trimmed); + return Ok(trimmed.to_string()); + } + } + + // Generate a random ID as a last resort + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + use std::time::{SystemTime, UNIX_EPOCH}; + + let mut hasher = DefaultHasher::new(); + if let Ok(n) = SystemTime::now().duration_since(UNIX_EPOCH) { + n.hash(&mut hasher); + } + let pid = std::process::id(); + pid.hash(&mut hasher); + let random_id = format!("cli_{:x}", hasher.finish()); + + // Save it for future use + if let Some(parent) = machine_id_path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let _ = std::fs::write(&machine_id_path, &random_id); + + Ok(random_id) +} diff --git a/crates/ov_cli/src/main.rs b/crates/ov_cli/src/main.rs index 2c276269..abc55547 100644 --- a/crates/ov_cli/src/main.rs +++ b/crates/ov_cli/src/main.rs @@ -331,9 +331,9 @@ enum Commands { /// Message to send to the agent #[arg(short, long)] message: Option, - /// Session ID - #[arg(short, long, default_value = "cli__chat__default")] - session: String, + /// Session ID (defaults to machine unique ID) + #[arg(short, long)] + session: Option, /// Render assistant output as Markdown #[arg(long = "markdown", default_value = "true")] markdown: bool, @@ -576,10 +576,11 @@ async fn main() { handle_tui(uri, ctx).await } Commands::Chat { message, session, markdown, logs: _ } => { - let cmd = commands::chat_v2::ChatCommand { + let session_id = session.or_else(|| config::get_or_create_machine_id().ok()); + let cmd = commands::chat::ChatCommand { endpoint: std::env::var("VIKINGBOT_ENDPOINT").unwrap_or_else(|_| "http://localhost:1933/bot/v1".to_string()), api_key: std::env::var("VIKINGBOT_API_KEY").ok(), - session: Some(session), + session: session_id, user: "cli_user".to_string(), message, stream: false, From 6eebae3576839ddc78913d479f52df2511664594 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 17:08:30 +0800 Subject: [PATCH 10/21] Remove unsupported --logs parameter from chat command --- crates/ov_cli/src/main.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/crates/ov_cli/src/main.rs b/crates/ov_cli/src/main.rs index abc55547..b2e046d7 100644 --- a/crates/ov_cli/src/main.rs +++ b/crates/ov_cli/src/main.rs @@ -337,9 +337,6 @@ enum Commands { /// Render assistant output as Markdown #[arg(long = "markdown", default_value = "true")] markdown: bool, - /// Show vikingbot runtime logs during chat - #[arg(long = "logs", default_value = "false")] - logs: bool, }, /// Configuration management Config { @@ -575,7 +572,7 @@ async fn main() { Commands::Tui { uri } => { handle_tui(uri, ctx).await } - Commands::Chat { message, session, markdown, logs: _ } => { + Commands::Chat { message, session, markdown } => { let session_id = session.or_else(|| config::get_or_create_machine_id().ok()); let cmd = commands::chat::ChatCommand { endpoint: std::env::var("VIKINGBOT_ENDPOINT").unwrap_or_else(|_| "http://localhost:1933/bot/v1".to_string()), From 21f236e723605248861f07634295ac5b880bb47b Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 17:33:05 +0800 Subject: [PATCH 11/21] =?UTF-8?q?=E7=BB=9F=E4=B8=80=20Python=20=E5=92=8C?= =?UTF-8?q?=20Rust=20CLI=20=E7=9A=84=E9=BB=98=E8=AE=A4=20session=20ID=20?= =?UTF-8?q?=E7=94=9F=E6=88=90=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot/scripts/test_all.sh | 50 +++++++++++ bot/tests/test_chat_commands.py | 133 ++++++++++++++++++++++++++++++ bot/uv.lock | 2 +- bot/vikingbot/channels/openapi.py | 2 +- bot/vikingbot/cli/commands.py | 78 +++++++++++++++--- openviking/server/bootstrap.py | 84 +++++++++++++++++-- 6 files changed, 328 insertions(+), 21 deletions(-) create mode 100755 bot/scripts/test_all.sh create mode 100644 bot/tests/test_chat_commands.py diff --git a/bot/scripts/test_all.sh b/bot/scripts/test_all.sh new file mode 100755 index 00000000..7c17cef1 --- /dev/null +++ b/bot/scripts/test_all.sh @@ -0,0 +1,50 @@ +#!/bin/bash + +# test_all.sh - Run all tests without triggering uv sync +# Usage: scripts/test_all.sh + +set -e + +# Get the directory of this script +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" + +echo "==========================================" +echo "Running OpenViking Tests" +echo "==========================================" +echo "" + +# Check if virtual environment exists +VENV_PYTHON="$PROJECT_ROOT/.venv/bin/python" +if [ ! -x "$VENV_PYTHON" ]; then + echo "Error: Virtual environment not found at $VENV_PYTHON" + echo "Please create a virtual environment first." + exit 1 +fi + +echo "Using Python from virtual environment: $VENV_PYTHON" +echo "" + +# Change to project root +cd "$PROJECT_ROOT" + +# Run pytest directly +echo "Running tests..." +echo "-----------------------------------------" +"$VENV_PYTHON" -m pytest tests/ -v +TEST_EXIT_CODE=$? +echo "-----------------------------------------" +echo "" + +# Show summary +if [ $TEST_EXIT_CODE -eq 0 ]; then + echo "==========================================" + echo "✅ All tests passed!" + echo "==========================================" +else + echo "==========================================" + echo "❌ Some tests failed (exit code: $TEST_EXIT_CODE)" + echo "==========================================" +fi + +exit $TEST_EXIT_CODE diff --git a/bot/tests/test_chat_commands.py b/bot/tests/test_chat_commands.py new file mode 100644 index 00000000..a88e5690 --- /dev/null +++ b/bot/tests/test_chat_commands.py @@ -0,0 +1,133 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for vikingbot chat commands""" + +import subprocess +import sys +from pathlib import Path + +import pytest + + +def test_vikingbot_chat_help(): + """Test that vikingbot chat --help shows correct options""" + # Run vikingbot chat --help + result = subprocess.run( + [sys.executable, "-m", "vikingbot.cli.commands", "chat", "--help"], + capture_output=True, + text=True, + ) + + # Check exit code + assert result.returncode == 0, f"Command failed: {result.stderr}" + + # Check that expected options are present + assert "--message" in result.stdout or "-m" in result.stdout + assert "--session" in result.stdout or "-s" in result.stdout + assert "--markdown" in result.stdout + assert "--no-markdown" in result.stdout + assert "--logs" in result.stdout + assert "--no-logs" in result.stdout + + +def test_vikingbot_command_exists(): + """Test that vikingbot command can be invoked""" + # Just check that the main module can be imported and shows help + result = subprocess.run( + [sys.executable, "-m", "vikingbot.cli.commands", "--help"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert "chat" in result.stdout + assert "Interact with the agent directly" in result.stdout + + +def test_vikingbot_gateway_help(): + """Test that gateway command help works""" + result = subprocess.run( + [sys.executable, "-m", "vikingbot.cli.commands", "gateway", "--help"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert "--port" in result.stdout or "-p" in result.stdout + + +def test_single_turn_channel_import(): + """Test that SingleTurnChannel can be imported""" + from vikingbot.channels.single_turn import SingleTurnChannel, SingleTurnChannelConfig + + assert SingleTurnChannel is not None + assert SingleTurnChannelConfig is not None + + +def test_chat_channel_import(): + """Test that ChatChannel can be imported""" + from vikingbot.channels.chat import ChatChannel, ChatChannelConfig + + assert ChatChannel is not None + assert ChatChannelConfig is not None + + +def test_prepare_agent_channel_function(): + """Test that prepare_agent_channel function exists and can be imported""" + from vikingbot.cli.commands import prepare_agent_channel + + assert prepare_agent_channel is not None + assert callable(prepare_agent_channel) + + +def test_chat_command_function_exists(): + """Test that the chat command function is registered""" + from vikingbot.cli.commands import app + + # Check that we can get the chat command info + # Typer stores commands differently, let's just verify the chat attribute exists + assert hasattr(app, "commands") or hasattr(app, "registered_commands") + + # Try calling the chat command with --help as a better verification + import subprocess + import sys + + result = subprocess.run( + [sys.executable, "-m", "vikingbot.cli.commands", "chat", "--help"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert "Interact with the agent directly" in result.stdout + + +@pytest.mark.integration +def test_chat_single_turn_dry_run(): + """ + Dry run test for single-turn chat (doesn't actually call LLM) + This tests the infrastructure without requiring API keys + """ + # This test just verifies the modules load correctly + # A full integration test would need API keys and a running agent + from vikingbot.channels.single_turn import SingleTurnChannel, SingleTurnChannelConfig + from vikingbot.bus.queue import MessageBus + + config = SingleTurnChannelConfig() + bus = MessageBus() + + # Just verify we can instantiate the channel without errors + channel = SingleTurnChannel( + config, + bus, + workspace_path=Path("/tmp"), + message="Hello, test", + session_id="test-session", + markdown=True, + ) + + assert channel is not None + assert channel.name == "single_turn" + assert channel.message == "Hello, test" + assert channel.session_id == "test-session" + diff --git a/bot/uv.lock b/bot/uv.lock index 638cb082..39c55752 100644 --- a/bot/uv.lock +++ b/bot/uv.lock @@ -3633,7 +3633,7 @@ wheels = [ [[package]] name = "vikingbot" -version = "0.1.1" +version = "0.1.2" source = { editable = "." } dependencies = [ { name = "beautifulsoup4" }, diff --git a/bot/vikingbot/channels/openapi.py b/bot/vikingbot/channels/openapi.py index 20afb1d1..74f8ee1e 100644 --- a/bot/vikingbot/channels/openapi.py +++ b/bot/vikingbot/channels/openapi.py @@ -37,7 +37,7 @@ class OpenAPIChannelConfig(BaseChannelConfig): """Configuration for OpenAPI channel.""" enabled: bool = True - type: str = "openapi" + type: str = "cli" api_key: str = "" # If empty, no auth required allow_from: list[str] = [] max_concurrent_requests: int = 100 diff --git a/bot/vikingbot/cli/commands.py b/bot/vikingbot/cli/commands.py index 31d41bdd..e297f45f 100644 --- a/bot/vikingbot/cli/commands.py +++ b/bot/vikingbot/cli/commands.py @@ -6,6 +6,8 @@ import select import signal import sys +import hashlib +import time from pathlib import Path import typer @@ -46,6 +48,60 @@ EXIT_COMMANDS = {"exit", "quit", "/exit", "/quit", ":q"} +def get_or_create_machine_id() -> str: + """Get or create a unique machine ID. + + This function will: + 1. Try to get the hostname from environment variables + 2. If that fails, generate a random ID and store it in the config directory + + The logic is kept consistent with the Rust implementation. + """ + home = Path.home() + machine_id_path = home / ".openviking" / "machine_id" + + # Try to read existing machine ID + if machine_id_path.exists(): + try: + content = machine_id_path.read_text().strip() + if content: + return content + except Exception: + pass + + # Try to get hostname from environment variables + for env_var in ("HOSTNAME", "COMPUTERNAME"): + hostname = os.environ.get(env_var) + if hostname: + trimmed = hostname.strip() + if trimmed: + # Save it for future use + machine_id_path.parent.mkdir(parents=True, exist_ok=True) + try: + machine_id_path.write_text(trimmed) + except Exception: + pass + return trimmed + + # Generate a random ID as a last resort + hasher = hashlib.sha256() + # Mix current timestamp + hasher.update(str(time.time_ns()).encode()) + # Mix process ID + hasher.update(str(os.getpid()).encode()) + # Take first 16 hex chars and prepend with "cli_" + random_id = f"cli_{hasher.hexdigest()[:16]}" + + # Save it for future use + machine_id_path.parent.mkdir(parents=True, exist_ok=True) + try: + machine_id_path.write_text(random_id) + except Exception: + pass + + return random_id + + def _init_bot_data(config): """Initialize bot data directory and set global paths.""" set_bot_data_path(config.bot_data_path) @@ -241,16 +297,16 @@ def gateway( async def run(): import uvicorn - # Setup OpenAPI routes before starting - openapi_channel = None - for name, channel in channels.channels.items(): - if hasattr(channel, 'name') and channel.name == "openapi": - openapi_channel = channel - break - - if openapi_channel is not None and hasattr(openapi_channel, '_setup_routes'): - openapi_channel._setup_routes() - logger.info("OpenAPI routes registered") + # # Setup OpenAPI routes before starting + # openapi_channel = None + # for name, channel in channels.channels.items(): + # if hasattr(channel, 'name') and channel.name == "openapi": + # openapi_channel = channel + # break + # + # if openapi_channel is not None and hasattr(openapi_channel, '_setup_routes'): + # openapi_channel._setup_routes() + # logger.info("OpenAPI routes registered") # Start uvicorn server for OpenAPI config_uvicorn = uvicorn.Config( @@ -546,7 +602,7 @@ def chat( is_single_turn = message is not None # Use unified default session ID if session_id is None: - session_id = "default" + session_id = get_or_create_machine_id() cron = prepare_cron(bus, quiet=is_single_turn) channels = prepare_agent_channel(config, bus, message, session_id, markdown, logs, eval) agent_loop = prepare_agent_loop(config, bus, session_manager, cron, quiet=is_single_turn, eval=eval) diff --git a/openviking/server/bootstrap.py b/openviking/server/bootstrap.py index 83b8149c..ee21cae4 100644 --- a/openviking/server/bootstrap.py +++ b/openviking/server/bootstrap.py @@ -3,6 +3,7 @@ """Bootstrap script for OpenViking HTTP Server.""" import argparse +import datetime import os import subprocess import sys @@ -70,6 +71,25 @@ def main(): dest="bot_url", help="Vikingbot OpenAPIChannel URL (default: http://localhost:18790)", ) + parser.add_argument( + "--enable-bot-logging", + action="store_true", + dest="enable_bot_logging", + default=None, + help="Enable logging vikingbot output to files (default: True when --with-bot is used)", + ) + parser.add_argument( + "--disable-bot-logging", + action="store_false", + dest="enable_bot_logging", + help="Disable logging vikingbot output to files", + ) + parser.add_argument( + "--bot-log-dir", + type=str, + default=os.path.expanduser("~/.openviking/data/bot/logs"), + help="Directory to store vikingbot log files", + ) args = parser.parse_args() @@ -100,10 +120,15 @@ def main(): if config.with_bot: print(f"Bot API proxy enabled, forwarding to {config.bot_api_url}") + # Determine if bot logging should be enabled + enable_bot_logging = args.enable_bot_logging + if enable_bot_logging is None: + enable_bot_logging = args.with_bot + # Start vikingbot gateway if --with-bot is set bot_process = None if args.with_bot: - bot_process = _start_vikingbot_gateway() + bot_process = _start_vikingbot_gateway(enable_bot_logging, args.bot_log_dir) try: uvicorn.run(app, host=config.host, port=config.port, log_config=None) @@ -113,7 +138,7 @@ def main(): _stop_vikingbot_gateway(bot_process) -def _start_vikingbot_gateway() -> subprocess.Popen: +def _start_vikingbot_gateway(enable_logging: bool, log_dir: str) -> subprocess.Popen: """Start vikingbot gateway as a subprocess.""" print("Starting vikingbot gateway...") @@ -140,6 +165,28 @@ def _start_vikingbot_gateway() -> subprocess.Popen: print(" cd bot && uv pip install -e '.[dev]'") return None + # Prepare logging + log_file = None + stdout_handler = subprocess.PIPE + stderr_handler = subprocess.PIPE + + if enable_logging: + try: + os.makedirs(log_dir, exist_ok=True) + log_filename = "vikingbot.log" + log_file_path = os.path.join(log_dir, log_filename) + log_file = open(log_file_path, "a") + stdout_handler = log_file + stderr_handler = log_file + print(f"Vikingbot logs will be written to: {log_file_path}") + except Exception as e: + print(f"Warning: Failed to setup bot logging: {e}") + if log_file: + log_file.close() + log_file = None + stdout_handler = subprocess.PIPE + stderr_handler = subprocess.PIPE + # Start vikingbot gateway process try: # Set environment to ensure it uses the same Python environment @@ -147,8 +194,8 @@ def _start_vikingbot_gateway() -> subprocess.Popen: process = subprocess.Popen( vikingbot_cmd, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, + stdout=stdout_handler, + stderr=stderr_handler, text=True, env=env, ) @@ -157,16 +204,30 @@ def _start_vikingbot_gateway() -> subprocess.Popen: time.sleep(2) if process.poll() is not None: # Process exited early - stdout, stderr = process.communicate(timeout=1) - print(f"Warning: vikingbot gateway exited early (code {process.returncode})") - if stderr: - print(f"Error: {stderr[:500]}") + if log_file: + log_file.close() + with open(log_file_path, "r") as f: + output = f.read() + print(f"Warning: vikingbot gateway exited early (code {process.returncode})") + if output: + print(f"Output: {output[:500]}") + else: + stdout, stderr = process.communicate(timeout=1) + print(f"Warning: vikingbot gateway exited early (code {process.returncode})") + if stderr: + print(f"Error: {stderr[:500]}") return None print(f"Vikingbot gateway started (PID: {process.pid})") + + # Store the log file with the process so we can close it later + process._log_file = log_file + return process except Exception as e: + if log_file: + log_file.close() print(f"Warning: Failed to start vikingbot gateway: {e}") return None @@ -191,6 +252,13 @@ def _stop_vikingbot_gateway(process: subprocess.Popen) -> None: print("Vikingbot gateway force killed.") except Exception as e: print(f"Error stopping vikingbot gateway: {e}") + finally: + # Close the log file if it exists + if hasattr(process, "_log_file") and process._log_file is not None: + try: + process._log_file.close() + except Exception as e: + print(f"Error closing bot log file: {e}") if __name__ == "__main__": From e2fc9efe9c32b58d65abbe0afd74e2605d31824d Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 17:42:52 +0800 Subject: [PATCH 12/21] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=97=A5=E5=BF=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot/vikingbot/cli/commands.py | 2 +- bot/vikingbot/integrations/langfuse.py | 2 +- bot/vikingbot/utils/tracing.py | 2 -- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/bot/vikingbot/cli/commands.py b/bot/vikingbot/cli/commands.py index e297f45f..96f2d5a4 100644 --- a/bot/vikingbot/cli/commands.py +++ b/bot/vikingbot/cli/commands.py @@ -336,7 +336,7 @@ def prepare_agent_loop(config, bus, session_manager, cron, quiet: bool = False, source_workspace_path = get_source_workspace_path() sandbox_manager = SandboxManager(config, sandbox_parent_path, source_workspace_path) if config.sandbox.backend == "direct": - logger.warning("Sandbox: disabled (using DIRECT mode - commands run directly on host)") + logger.warning("[SANDBOX] disabled (using DIRECT mode - commands run directly on host)") else: logger.info(f"Sandbox: enabled (backend={config.sandbox.backend}, mode={config.sandbox.mode})") diff --git a/bot/vikingbot/integrations/langfuse.py b/bot/vikingbot/integrations/langfuse.py index cd6f8d2c..c6ea7ad2 100644 --- a/bot/vikingbot/integrations/langfuse.py +++ b/bot/vikingbot/integrations/langfuse.py @@ -61,7 +61,7 @@ def __init__( def get_instance(cls) -> "LangfuseClient": """Get the singleton instance.""" if cls._instance is None: - logger.warning("[LANGFUSE] No instance set, creating default (disabled) instance") + logger.warning("[LANGFUSE] disabled") cls._instance = LangfuseClient(enabled=False) return cls._instance diff --git a/bot/vikingbot/utils/tracing.py b/bot/vikingbot/utils/tracing.py index 9dee8534..58039a08 100644 --- a/bot/vikingbot/utils/tracing.py +++ b/bot/vikingbot/utils/tracing.py @@ -161,8 +161,6 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> T: with langfuse.propagate_attributes(session_id=session_id, user_id=user_id): return await wrapped_func(*args, **kwargs) else: - if not langfuse.enabled: - logger.warning(f"[LANGFUSE] Client not enabled") if not has_propagate: logger.warning(f"[LANGFUSE] propagate_attributes not available") return await wrapped_func(*args, **kwargs) From b92e9284861919316f9541f14d433b6c1b774f9b Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 18:49:40 +0800 Subject: [PATCH 13/21] =?UTF-8?q?=E5=8E=BB=E6=8E=89log=E4=BE=9D=E8=B5=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- openviking/server/routers/bot.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/openviking/server/routers/bot.py b/openviking/server/routers/bot.py index b39f16f8..bc819a18 100644 --- a/openviking/server/routers/bot.py +++ b/openviking/server/routers/bot.py @@ -11,10 +11,13 @@ import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from fastapi.responses import JSONResponse, StreamingResponse -from loguru import logger + +from openviking_cli.utils.logger import get_logger router = APIRouter(prefix="", tags=["bot"]) +logger = get_logger(__name__) + # Bot API configuration - set when --with-bot is enabled BOT_API_URL: Optional[str] = None # e.g., "http://localhost:18791" From 7ceb2d2296b268d11ae0d79a4409f6f0152fc251 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 20:28:47 +0800 Subject: [PATCH 14/21] docs: add VikingBot quick start section to READMEs --- README.md | 15 +++++++++++++++ README_CN.md | 15 +++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/README.md b/README.md index 04e975cc..3d696d51 100644 --- a/README.md +++ b/README.md @@ -437,6 +437,21 @@ ov grep "openviking" --uri viking://resources/volcengine/OpenViking/docs/zh Congratulations! You have successfully run OpenViking 🎉 +### VikingBot Quick Start + +VikingBot is an AI agent framework built on top of OpenViking. Here's how to get started: + +```bash +# Install VikingBot from source (in OpenViking root directory) +uv pip install -e bot/ + +# Start OpenViking server with Bot enabled +openviking-server --with-bot + +# In another terminal, start interactive chat +ov chat +``` + --- ## Server Deployment Details diff --git a/README_CN.md b/README_CN.md index 8ac5f417..2ddc20ad 100644 --- a/README_CN.md +++ b/README_CN.md @@ -436,6 +436,21 @@ ov grep "openviking" --uri viking://resources/volcengine/OpenViking/docs/zh 恭喜！您已成功运行 OpenViking 🎉 +### VikingBot 快速开始 + +VikingBot 是构建在 OpenViking 之上的 AI 智能体框架。以下是快速开始指南： + +```bash +# 在 OpenViking 源码根目录下安装 VikingBot +uv pip install -e bot/ + +# 启动 OpenViking 服务器（同时启动 Bot） +openviking-server --with-bot + +# 在另一个终端启动交互式聊天 +ov chat +``` + --- ## 服务器部署详情 From 59f4e87ba04b73093cc749e477c1df27d1284d24 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 20:45:48 +0800 Subject: [PATCH 15/21] fix: use vikingbot chat instead of ov chat in READMEs --- README.md | 2 +- README_CN.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 3d696d51..58d66d41 100644 --- a/README.md +++ b/README.md @@ -449,7 +449,7 @@ uv pip install -e bot/ openviking-server --with-bot # In another terminal, start interactive chat -ov chat +vikingbot chat ``` --- diff --git a/README_CN.md b/README_CN.md index 2ddc20ad..2e28435d 100644 --- a/README_CN.md +++ b/README_CN.md @@ -448,7 +448,7 @@ uv pip install -e bot/ openviking-server --with-bot # 在另一个终端启动交互式聊天 -ov chat +vikingbot chat ``` --- From 8fb714af598f9a142021e56a4305af19f3b569b4 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 20:47:35 +0800 Subject: [PATCH 16/21] Revert "fix: use vikingbot chat instead of ov chat in READMEs" This reverts commit 59f4e87ba04b73093cc749e477c1df27d1284d24. --- README.md | 2 +- README_CN.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 58d66d41..3d696d51 100644 --- a/README.md +++ b/README.md @@ -449,7 +449,7 @@ uv pip install -e bot/ openviking-server --with-bot # In another terminal, start interactive chat -vikingbot chat +ov chat ``` --- diff --git a/README_CN.md b/README_CN.md index 2e28435d..2ddc20ad 100644 --- a/README_CN.md +++ b/README_CN.md @@ -448,7 +448,7 @@ uv pip install -e bot/ openviking-server --with-bot # 在另一个终端启动交互式聊天 -vikingbot chat +ov chat ``` --- From fe82005dd584b7901ce1b616bfa7b31f1b1239ae Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 21:07:40 +0800 Subject: [PATCH 17/21] fix: use UUID v4 for machine ID generation in both Rust and Python --- bot/vikingbot/cli/commands.py | 44 ++++++----------------------------- crates/ov_cli/Cargo.toml | 1 + crates/ov_cli/src/config.rs | 35 +++++----------------------- crates/ov_cli/src/utils.rs | 19 +++++++++++++++ 4 files changed, 33 insertions(+), 66 deletions(-) create mode 100644 crates/ov_cli/src/utils.rs diff --git a/bot/vikingbot/cli/commands.py b/bot/vikingbot/cli/commands.py index 96f2d5a4..df904316 100644 --- a/bot/vikingbot/cli/commands.py +++ b/bot/vikingbot/cli/commands.py @@ -52,11 +52,12 @@ def get_or_create_machine_id() -> str: """Get or create a unique machine ID. This function will: - 1. Try to get the hostname from environment variables - 2. If that fails, generate a random ID and store it in the config directory + 1. Try to read an existing machine ID from the config directory + 2. If not found, generate a new UUID v4 and store it The logic is kept consistent with the Rust implementation. """ + import uuid home = Path.home() machine_id_path = home / ".openviking" / "machine_id" @@ -69,37 +70,17 @@ def get_or_create_machine_id() -> str: except Exception: pass - # Try to get hostname from environment variables - for env_var in ("HOSTNAME", "COMPUTERNAME"): - hostname = os.environ.get(env_var) - if hostname: - trimmed = hostname.strip() - if trimmed: - # Save it for future use - machine_id_path.parent.mkdir(parents=True, exist_ok=True) - try: - machine_id_path.write_text(trimmed) - except Exception: - pass - return trimmed - - # Generate a random ID as a last resort - hasher = hashlib.sha256() - # Mix current timestamp - hasher.update(str(time.time_ns()).encode()) - # Mix process ID - hasher.update(str(os.getpid()).encode()) - # Take first 16 hex chars and prepend with "cli_" - random_id = f"cli_{hasher.hexdigest()[:16]}" + # Generate a new UUID v4 + new_id = str(uuid.uuid4()) # Save it for future use machine_id_path.parent.mkdir(parents=True, exist_ok=True) try: - machine_id_path.write_text(random_id) + machine_id_path.write_text(new_id) except Exception: pass - return random_id + return new_id def _init_bot_data(config): @@ -297,17 +278,6 @@ def gateway( async def run(): import uvicorn - # # Setup OpenAPI routes before starting - # openapi_channel = None - # for name, channel in channels.channels.items(): - # if hasattr(channel, 'name') and channel.name == "openapi": - # openapi_channel = channel - # break - # - # if openapi_channel is not None and hasattr(openapi_channel, '_setup_routes'): - # openapi_channel._setup_routes() - # logger.info("OpenAPI routes registered") - # Start uvicorn server for OpenAPI config_uvicorn = uvicorn.Config( fastapi_app, diff --git a/crates/ov_cli/Cargo.toml b/crates/ov_cli/Cargo.toml index 08b74f03..671b160e 100644 --- a/crates/ov_cli/Cargo.toml +++ b/crates/ov_cli/Cargo.toml @@ -30,3 +30,4 @@ tempfile = "3.12" url = "2.5" walkdir = "2.5" rustyline = "14.0" +uuid = { version = "1.0", features = ["v4", "serde"] } diff --git a/crates/ov_cli/src/config.rs b/crates/ov_cli/src/config.rs index f7ba3a1c..8da8000b 100644 --- a/crates/ov_cli/src/config.rs +++ b/crates/ov_cli/src/config.rs @@ -102,8 +102,8 @@ pub fn default_config_path() -> Result { /// Get or create a unique machine ID. /// /// This function will: -/// 1. Try to get the hostname from environment variables -/// 2. If that fails, generate a random ID and store it in the config directory +/// 1. Try to read an existing machine ID from the config directory +/// 2. If not found, generate a new UUID v4 and store it pub fn get_or_create_machine_id() -> Result { let home = dirs::home_dir() .ok_or_else(|| Error::Config("Could not determine home directory".to_string()))?; @@ -119,37 +119,14 @@ pub fn get_or_create_machine_id() -> Result { } } - // Try to get hostname from environment variables - if let Ok(hostname) = std::env::var("HOSTNAME").or_else(|_| std::env::var("COMPUTERNAME")) { - let trimmed = hostname.trim(); - if !trimmed.is_empty() { - // Save it for future use - if let Some(parent) = machine_id_path.parent() { - let _ = std::fs::create_dir_all(parent); - } - let _ = std::fs::write(&machine_id_path, trimmed); - return Ok(trimmed.to_string()); - } - } - - // Generate a random ID as a last resort - use std::collections::hash_map::DefaultHasher; - use std::hash::{Hash, Hasher}; - use std::time::{SystemTime, UNIX_EPOCH}; - - let mut hasher = DefaultHasher::new(); - if let Ok(n) = SystemTime::now().duration_since(UNIX_EPOCH) { - n.hash(&mut hasher); - } - let pid = std::process::id(); - pid.hash(&mut hasher); - let random_id = format!("cli_{:x}", hasher.finish()); + // Generate a new UUID v4 + let new_id = uuid::Uuid::new_v4().to_string(); // Save it for future use if let Some(parent) = machine_id_path.parent() { let _ = std::fs::create_dir_all(parent); } - let _ = std::fs::write(&machine_id_path, &random_id); + let _ = std::fs::write(&machine_id_path, &new_id); - Ok(random_id) + Ok(new_id) } diff --git a/crates/ov_cli/src/utils.rs b/crates/ov_cli/src/utils.rs new file mode 100644 index 00000000..24dee374 --- /dev/null +++ b/crates/ov_cli/src/utils.rs @@ -0,0 +1,19 @@ +//! Utility functions used across the crate. + +/// Safely truncate a string at a UTF-8 character boundary +pub fn truncate_utf8(s: &str, max_bytes: usize) -> &str { + if s.len() <= max_bytes { + return s; + } + + let mut boundary = max_bytes; + while boundary > 0 && !s.is_char_boundary(boundary) { + boundary -= 1; + } + + if boundary == 0 { + "" + } else { + &s[..boundary] + } +} From 86ed2e17bb0f2d4b152e047103a1c4af640d9517 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 21:11:10 +0800 Subject: [PATCH 18/21] refactor: move truncate_utf8 to utils, fix chat history path, and use BotProcess dataclass --- crates/ov_cli/src/commands/chat.rs | 22 +----------- crates/ov_cli/src/commands/chat_v2.rs | 38 +++++++-------------- crates/ov_cli/src/main.rs | 1 + openviking/server/bootstrap.py | 49 +++++++++++++++------------ 4 files changed, 43 insertions(+), 67 deletions(-) diff --git a/crates/ov_cli/src/commands/chat.rs b/crates/ov_cli/src/commands/chat.rs index 782b540d..3fb743f4 100644 --- a/crates/ov_cli/src/commands/chat.rs +++ b/crates/ov_cli/src/commands/chat.rs @@ -3,28 +3,8 @@ use std::io::{BufRead, Write}; use std::time::Duration; -/// Safely truncate a string at a UTF-8 character boundary -fn truncate_utf8(s: &str, max_bytes: usize) -> &str { - if s.len() <= max_bytes { - return s; - } - - // Find the last valid UTF-8 character boundary before or at max_bytes - let mut boundary = max_bytes; - while boundary > 0 && !s.is_char_boundary(boundary) { - boundary -= 1; - } - - // If we couldn't find a boundary (unlikely), just return empty string - // Otherwise return up to the boundary - if boundary == 0 { - "" - } else { - &s[..boundary] - } -} - use clap::Parser; +use crate::utils::truncate_utf8; use reqwest::Client; use serde::{Deserialize, Serialize}; diff --git a/crates/ov_cli/src/commands/chat_v2.rs b/crates/ov_cli/src/commands/chat_v2.rs index afafcb6e..966603c3 100644 --- a/crates/ov_cli/src/commands/chat_v2.rs +++ b/crates/ov_cli/src/commands/chat_v2.rs @@ -1,6 +1,6 @@ - //! Chat command for interacting with Vikingbot via OpenAPI (v2 with rustyline) +use std::path::PathBuf; use std::time::Duration; use clap::Parser; @@ -10,27 +10,10 @@ use rustyline::error::ReadlineError; use rustyline::{Editor, history::FileHistory}; use crate::error::{Error, Result}; +use crate::utils::truncate_utf8; const DEFAULT_ENDPOINT: &str = "http://localhost:1933/bot/v1"; -/// Safely truncate a string at a UTF-8 character boundary -fn truncate_utf8(s: &str, max_bytes: usize) -> &str { - if s.len() <= max_bytes { - return s; - } - - let mut boundary = max_bytes; - while boundary > 0 && !s.is_char_boundary(boundary) { - boundary -= 1; - } - - if boundary == 0 { - "" - } else { - &s[..boundary] - } -} - #[derive(Debug, Parser)] pub struct ChatCommand { #[arg(short, long, default_value = DEFAULT_ENDPOINT)] @@ -195,7 +178,8 @@ impl ChatCommand { let mut rl: Editor<(), FileHistory> = Editor::new() .map_err(|e| Error::Client(format!("Failed to create line editor: {}", e)))?; - let _ = rl.load_history("ov_chat_history.txt"); + let history_path = get_chat_history_path()?; + let _ = rl.load_history(&history_path); loop { let readline = rl.readline("\x1b[1;32mYou:\x1b[0m "); @@ -205,7 +189,7 @@ impl ChatCommand { let trimmed = line.trim().to_string(); if !trimmed.is_empty() { let _ = rl.add_history_entry(line); - let _ = rl.save_history("ov_chat_history.txt"); + let _ = rl.save_history(&history_path); } trimmed } @@ -313,15 +297,11 @@ impl ChatCommand { println!("\nGoodbye!"); Ok(()) } -} -impl ChatCommand { pub async fn run(&self) -> Result<()> { self.execute().await } -} -impl ChatCommand { #[allow(clippy::too_many_arguments)] pub fn new( endpoint: String, @@ -344,3 +324,11 @@ impl ChatCommand { } } +fn get_chat_history_path() -> Result { + let home = dirs::home_dir() + .ok_or_else(|| Error::Config("Could not determine home directory".to_string()))?; + let config_dir = home.join(".openviking"); + std::fs::create_dir_all(&config_dir) + .map_err(|e| Error::Config(format!("Failed to create config directory: {}", e)))?; + Ok(config_dir.join("ov_chat_history.txt")) +} diff --git a/crates/ov_cli/src/main.rs b/crates/ov_cli/src/main.rs index b2e046d7..f41d2081 100644 --- a/crates/ov_cli/src/main.rs +++ b/crates/ov_cli/src/main.rs @@ -4,6 +4,7 @@ mod config; mod error; mod output; mod tui; +mod utils; use clap::{Parser, Subcommand}; use config::Config; diff --git a/openviking/server/bootstrap.py b/openviking/server/bootstrap.py index ee21cae4..942c3c91 100644 --- a/openviking/server/bootstrap.py +++ b/openviking/server/bootstrap.py @@ -3,7 +3,6 @@ """Bootstrap script for OpenViking HTTP Server.""" import argparse -import datetime import os import subprocess import sys @@ -11,11 +10,20 @@ import uvicorn +from dataclasses import dataclass +from typing import Optional + from openviking.server.app import create_app from openviking.server.config import load_server_config from openviking_cli.utils.logger import configure_uvicorn_logging +@dataclass +class BotProcess: + process: subprocess.Popen + log_file: Optional[object] = None + + def _get_version() -> str: try: from openviking import __version__ @@ -126,7 +134,7 @@ def main(): enable_bot_logging = args.with_bot # Start vikingbot gateway if --with-bot is set - bot_process = None + bot_process: Optional[BotProcess] = None if args.with_bot: bot_process = _start_vikingbot_gateway(enable_bot_logging, args.bot_log_dir) @@ -138,7 +146,7 @@ def main(): _stop_vikingbot_gateway(bot_process) -def _start_vikingbot_gateway(enable_logging: bool, log_dir: str) -> subprocess.Popen: +def _start_vikingbot_gateway(enable_logging: bool, log_dir: str) -> Optional[BotProcess]: """Start vikingbot gateway as a subprocess.""" print("Starting vikingbot gateway...") @@ -169,6 +177,7 @@ def _start_vikingbot_gateway(enable_logging: bool, log_dir: str) -> subprocess.P log_file = None stdout_handler = subprocess.PIPE stderr_handler = subprocess.PIPE + log_file_path = None if enable_logging: try: @@ -206,11 +215,12 @@ def _start_vikingbot_gateway(enable_logging: bool, log_dir: str) -> subprocess.P # Process exited early if log_file: log_file.close() - with open(log_file_path, "r") as f: - output = f.read() - print(f"Warning: vikingbot gateway exited early (code {process.returncode})") - if output: - print(f"Output: {output[:500]}") + if log_file_path: + with open(log_file_path, "r") as f: + output = f.read() + print(f"Warning: vikingbot gateway exited early (code {process.returncode})") + if output: + print(f"Output: {output[:500]}") else: stdout, stderr = process.communicate(timeout=1) print(f"Warning: vikingbot gateway exited early (code {process.returncode})") @@ -220,10 +230,7 @@ def _start_vikingbot_gateway(enable_logging: bool, log_dir: str) -> subprocess.P print(f"Vikingbot gateway started (PID: {process.pid})") - # Store the log file with the process so we can close it later - process._log_file = log_file - - return process + return BotProcess(process=process, log_file=log_file) except Exception as e: if log_file: @@ -232,31 +239,31 @@ def _start_vikingbot_gateway(enable_logging: bool, log_dir: str) -> subprocess.P return None -def _stop_vikingbot_gateway(process: subprocess.Popen) -> None: +def _stop_vikingbot_gateway(bot_process: BotProcess) -> None: """Stop the vikingbot gateway subprocess.""" - if process is None: + if bot_process is None: return - print(f"\nStopping vikingbot gateway (PID: {process.pid})...") + print(f"\nStopping vikingbot gateway (PID: {bot_process.process.pid})...") try: # Try graceful termination first - process.terminate() + bot_process.process.terminate() try: - process.wait(timeout=5) + bot_process.process.wait(timeout=5) print("Vikingbot gateway stopped gracefully.") except subprocess.TimeoutExpired: # Force kill if it doesn't stop in time - process.kill() - process.wait() + bot_process.process.kill() + bot_process.process.wait() print("Vikingbot gateway force killed.") except Exception as e: print(f"Error stopping vikingbot gateway: {e}") finally: # Close the log file if it exists - if hasattr(process, "_log_file") and process._log_file is not None: + if bot_process.log_file is not None: try: - process._log_file.close() + bot_process.log_file.close() except Exception as e: print(f"Error closing bot log file: {e}") From 824d15da9dafa768c6e23fa10e5e0d1d2b92c044 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 22:08:48 +0800 Subject: [PATCH 19/21] refactor: update machine ID generation and remove unused chat_v2 - Update Python to use py-machineid library - Update Rust to use machine-uid crate - Remove unused chat_v2.rs - Move machine ID from file storage to system-provided IDs - Add fallback to "default" if system ID is unavailable --- bot/pyproject.toml | 1 + bot/vikingbot/cli/commands.py | 36 +-- crates/ov_cli/Cargo.toml | 1 + crates/ov_cli/src/commands/chat.rs | 6 +- crates/ov_cli/src/commands/chat_v2.rs | 334 -------------------------- crates/ov_cli/src/commands/mod.rs | 1 - crates/ov_cli/src/config.rs | 32 +-- 7 files changed, 20 insertions(+), 391 deletions(-) delete mode 100644 crates/ov_cli/src/commands/chat_v2.rs diff --git a/bot/pyproject.toml b/bot/pyproject.toml index dc4ebbaa..af0c38c6 100644 --- a/bot/pyproject.toml +++ b/bot/pyproject.toml @@ -41,6 +41,7 @@ dependencies = [ "openviking>=0.1.18", "ddgs>=9.0.0", "gradio>=6.6.0", + "py-machineid>=1.0.0", ] [project.optional-dependencies] diff --git a/bot/vikingbot/cli/commands.py b/bot/vikingbot/cli/commands.py index df904316..98c2756e 100644 --- a/bot/vikingbot/cli/commands.py +++ b/bot/vikingbot/cli/commands.py @@ -49,38 +49,22 @@ def get_or_create_machine_id() -> str: - """Get or create a unique machine ID. + """Get a unique machine ID using py-machineid. - This function will: - 1. Try to read an existing machine ID from the config directory - 2. If not found, generate a new UUID v4 and store it - - The logic is kept consistent with the Rust implementation. + Uses the system's machine ID, falls back to "default" if unavailable. """ - import uuid - home = Path.home() - machine_id_path = home / ".openviking" / "machine_id" - - # Try to read existing machine ID - if machine_id_path.exists(): - try: - content = machine_id_path.read_text().strip() - if content: - return content - except Exception: - pass - - # Generate a new UUID v4 - new_id = str(uuid.uuid4()) - - # Save it for future use - machine_id_path.parent.mkdir(parents=True, exist_ok=True) try: - machine_id_path.write_text(new_id) + from machineid import machine_id + return machine_id() + except ImportError: + # Fallback if py-machineid is not installed + pass except Exception: pass - return new_id + # Default fallback + return "default" + def _init_bot_data(config): diff --git a/crates/ov_cli/Cargo.toml b/crates/ov_cli/Cargo.toml index 671b160e..d2ebcbc6 100644 --- a/crates/ov_cli/Cargo.toml +++ b/crates/ov_cli/Cargo.toml @@ -31,3 +31,4 @@ url = "2.5" walkdir = "2.5" rustyline = "14.0" uuid = { version = "1.0", features = ["v4", "serde"] } +machine-uid = "0.5" diff --git a/crates/ov_cli/src/commands/chat.rs b/crates/ov_cli/src/commands/chat.rs index 3fb743f4..cf985089 100644 --- a/crates/ov_cli/src/commands/chat.rs +++ b/crates/ov_cli/src/commands/chat.rs @@ -142,13 +142,13 @@ impl ChatCommand { "reasoning" => { let content = data.as_str().unwrap_or(""); if !self.no_format { - println!("\t\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); + println!(" \x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); } } "tool_call" => { let content = data.as_str().unwrap_or(""); if !self.no_format { - println!("\t\x1b[2m├─ Calling: {}\x1b[0m", content); + println!(" \x1b[2m├─ Calling: {}\x1b[0m", content); } } "tool_result" => { @@ -159,7 +159,7 @@ impl ChatCommand { } else { content.to_string() }; - println!("\t\x1b[2m└─ Result: {}\x1b[0m", truncated); + println!(" \x1b[2m└─ Result: {}\x1b[0m", truncated); } } _ => {} diff --git a/crates/ov_cli/src/commands/chat_v2.rs b/crates/ov_cli/src/commands/chat_v2.rs deleted file mode 100644 index 966603c3..00000000 --- a/crates/ov_cli/src/commands/chat_v2.rs +++ /dev/null @@ -1,334 +0,0 @@ -//! Chat command for interacting with Vikingbot via OpenAPI (v2 with rustyline) - -use std::path::PathBuf; -use std::time::Duration; - -use clap::Parser; -use reqwest::Client; -use serde::{Deserialize, Serialize}; -use rustyline::error::ReadlineError; -use rustyline::{Editor, history::FileHistory}; - -use crate::error::{Error, Result}; -use crate::utils::truncate_utf8; - -const DEFAULT_ENDPOINT: &str = "http://localhost:1933/bot/v1"; - -#[derive(Debug, Parser)] -pub struct ChatCommand { - #[arg(short, long, default_value = DEFAULT_ENDPOINT)] - pub endpoint: String, - - #[arg(short, long, env = "VIKINGBOT_API_KEY")] - pub api_key: Option, - - #[arg(short, long)] - pub session: Option, - - #[arg(short, long, default_value = "cli_user")] - pub user: String, - - #[arg(short = 'M', long)] - pub message: Option, - - #[arg(long)] - pub stream: bool, - - #[arg(long)] - pub no_format: bool, -} - -#[derive(Debug, Serialize, Deserialize)] -struct ChatMessage { - role: String, - content: String, -} - -#[derive(Debug, Serialize)] -struct ChatRequest { - message: String, - #[serde(skip_serializing_if = "Option::is_none")] - session_id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - user_id: Option, - stream: bool, - #[serde(skip_serializing_if = "Option::is_none")] - context: Option>, -} - -#[derive(Debug, Deserialize)] -struct ChatResponse { - session_id: String, - message: String, - #[serde(default)] - events: Option>, -} - -#[derive(Debug, Deserialize)] -struct StreamEvent { - event: String, - data: serde_json::Value, -} - -impl ChatCommand { - pub async fn execute(&self) -> Result<()> { - let client = Client::builder() - .timeout(Duration::from_secs(300)) - .build() - .map_err(|e| Error::Network(format!("Failed to create HTTP client: {}", e)))?; - - if let Some(message) = &self.message { - self.send_message(&client, message).await - } else { - self.run_interactive(&client).await - } - } - - async fn send_message(&self, client: &Client, message: &str) -> Result<()> { - let url = format!("{}/chat", self.endpoint); - - let request = ChatRequest { - message: message.to_string(), - session_id: self.session.clone(), - user_id: Some(self.user.clone()), - stream: false, - context: None, - }; - - let mut req_builder = client.post(&url).json(&request); - - if let Some(api_key) = &self.api_key { - req_builder = req_builder.header("X-API-Key", api_key); - } - - let response = req_builder - .send() - .await - .map_err(|e| Error::Network(format!("Failed to send request: {}", e)))?; - - if !response.status().is_success() { - let status = response.status(); - let text = response.text().await.unwrap_or_default(); - return Err(Error::Api(format!("Request failed ({}): {}", status, text))); - } - - let chat_response: ChatResponse = response - .json() - .await - .map_err(|e| Error::Parse(format!("Failed to parse response: {}", e)))?; - - if let Some(events) = &chat_response.events { - for event in events { - if let (Some(etype), Some(data)) = ( - event.get("type").and_then(|v| v.as_str()), - event.get("data"), - ) { - match etype { - "reasoning" => { - let content = data.as_str().unwrap_or(""); - if !self.no_format { - println!("\t\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); - } - } - "tool_call" => { - let content = data.as_str().unwrap_or(""); - if !self.no_format { - println!("\t\x1b[2m├─ Calling: {}\x1b[0m", content); - } - } - "tool_result" => { - let content = data.as_str().unwrap_or(""); - if !self.no_format { - let truncated = if content.len() > 150 { - format!("{}...", truncate_utf8(content, 150)) - } else { - content.to_string() - }; - println!("\t\x1b[2m└─ Result: {}\x1b[0m", truncated); - } - } - _ => {} - } - } - } - } - - if !self.no_format { - println!("\n\x1b[1;31mBot:\x1b[0m"); - println!("{}", chat_response.message); - println!(); - } else { - println!("{}", chat_response.message); - } - - Ok(()) - } - - async fn run_interactive(&self, client: &Client) -> Result<()> { - println!("Vikingbot Chat - Interactive Mode (v2)"); - println!("Endpoint: {}", self.endpoint); - if let Some(session) = &self.session { - println!("Session: {}", session); - } - println!("Type 'exit', 'quit', or press Ctrl+C to exit"); - println!("----------------------------------------\n"); - - let mut session_id = self.session.clone(); - - let mut rl: Editor<(), FileHistory> = Editor::new() - .map_err(|e| Error::Client(format!("Failed to create line editor: {}", e)))?; - - let history_path = get_chat_history_path()?; - let _ = rl.load_history(&history_path); - - loop { - let readline = rl.readline("\x1b[1;32mYou:\x1b[0m "); - - let input = match readline { - Ok(line) => { - let trimmed = line.trim().to_string(); - if !trimmed.is_empty() { - let _ = rl.add_history_entry(line); - let _ = rl.save_history(&history_path); - } - trimmed - } - Err(ReadlineError::Interrupted) => { - println!("\nGoodbye!"); - break; - } - Err(ReadlineError::Eof) => { - println!("\nGoodbye!"); - break; - } - Err(err) => { - eprintln!("\x1b[1;31mError reading input: {}\x1b[0m", err); - continue; - } - }; - - if input.is_empty() { - continue; - } - - if input.eq_ignore_ascii_case("exit") || input.eq_ignore_ascii_case("quit") { - println!("\nGoodbye!"); - break; - } - - let url = format!("{}/chat", self.endpoint); - - let request = ChatRequest { - message: input.to_string(), - session_id: session_id.clone(), - user_id: Some(self.user.clone()), - stream: false, - context: None, - }; - - let mut req_builder = client.post(&url).json(&request); - - if let Some(api_key) = &self.api_key { - req_builder = req_builder.header("X-API-Key", api_key); - } - - match req_builder.send().await { - Ok(response) => { - if response.status().is_success() { - match response.json::().await { - Ok(chat_response) => { - if session_id.is_none() { - session_id = Some(chat_response.session_id.clone()); - } - - if let Some(events) = chat_response.events { - for event in events { - if let (Some(etype), Some(data)) = ( - event.get("type").and_then(|v| v.as_str()), - event.get("data"), - ) { - match etype { - "reasoning" => { - let content = data.as_str().unwrap_or(""); - if content.len() > 100 { - println!("\t\x1b[2mThink: {}...\x1b[0m", truncate_utf8(content, 100)); - } else { - println!("\t\x1b[2mThink: {}\x1b[0m", content); - } - } - "tool_call" => { - println!("\t\x1b[2m├─ Calling: {}\x1b[0m", data.as_str().unwrap_or("")); - } - "tool_result" => { - let content = data.as_str().unwrap_or(""); - let truncated = if content.len() > 150 { - format!("{}...", truncate_utf8(content, 150)) - } else { - content.to_string() - }; - println!("\t\x1b[2m└─ Result: {}\x1b[0m", truncated); - } - _ => {} - } - } - } - } - - println!("\n\x1b[1;31mBot:\x1b[0m"); - println!("{}", chat_response.message); - println!(); - } - Err(e) => { - eprintln!("\x1b[1;31mError parsing response: {}\x1b[0m", e); - } - } - } else { - let status = response.status(); - let text = response.text().await.unwrap_or_default(); - eprintln!("\x1b[1;31mRequest failed ({}): {}\x1b[0m", status, text); - } - } - Err(e) => { - eprintln!("\x1b[1;31mFailed to send request: {}\x1b[0m", e); - } - } - } - - println!("\nGoodbye!"); - Ok(()) - } - - pub async fn run(&self) -> Result<()> { - self.execute().await - } - - #[allow(clippy::too_many_arguments)] - pub fn new( - endpoint: String, - api_key: Option, - session: Option, - user: String, - message: Option, - stream: bool, - no_format: bool, - ) -> Self { - Self { - endpoint, - api_key, - session, - user, - message, - stream, - no_format, - } - } -} - -fn get_chat_history_path() -> Result { - let home = dirs::home_dir() - .ok_or_else(|| Error::Config("Could not determine home directory".to_string()))?; - let config_dir = home.join(".openviking"); - std::fs::create_dir_all(&config_dir) - .map_err(|e| Error::Config(format!("Failed to create config directory: {}", e)))?; - Ok(config_dir.join("ov_chat_history.txt")) -} diff --git a/crates/ov_cli/src/commands/mod.rs b/crates/ov_cli/src/commands/mod.rs index 270766a6..1e8d1bcd 100644 --- a/crates/ov_cli/src/commands/mod.rs +++ b/crates/ov_cli/src/commands/mod.rs @@ -1,6 +1,5 @@ pub mod admin; pub mod chat; -pub mod chat_v2; pub mod content; pub mod search; pub mod filesystem; diff --git a/crates/ov_cli/src/config.rs b/crates/ov_cli/src/config.rs index 8da8000b..e27b3a65 100644 --- a/crates/ov_cli/src/config.rs +++ b/crates/ov_cli/src/config.rs @@ -99,34 +99,12 @@ pub fn default_config_path() -> Result { Ok(home.join(".openviking").join("ovcli.conf")) } -/// Get or create a unique machine ID. +/// Get a unique machine ID using machine-uid crate. /// -/// This function will: -/// 1. Try to read an existing machine ID from the config directory -/// 2. If not found, generate a new UUID v4 and store it +/// Uses the system's machine ID, falls back to "default" if unavailable. pub fn get_or_create_machine_id() -> Result { - let home = dirs::home_dir() - .ok_or_else(|| Error::Config("Could not determine home directory".to_string()))?; - let machine_id_path = home.join(".openviking").join("machine_id"); - - // Try to read existing machine ID - if machine_id_path.exists() { - if let Ok(content) = std::fs::read_to_string(&machine_id_path) { - let trimmed = content.trim(); - if !trimmed.is_empty() { - return Ok(trimmed.to_string()); - } - } + match machine_uid::get() { + Ok(id) => Ok(id), + Err(_) => Ok("default".to_string()), } - - // Generate a new UUID v4 - let new_id = uuid::Uuid::new_v4().to_string(); - - // Save it for future use - if let Some(parent) = machine_id_path.parent() { - let _ = std::fs::create_dir_all(parent); - } - let _ = std::fs::write(&machine_id_path, &new_id); - - Ok(new_id) } From c1e939807572a6b761b3d5562f28a4f60d8fa1dd Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Thu, 5 Mar 2026 22:15:09 +0800 Subject: [PATCH 20/21] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=A0=BC=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot/pyproject.toml | 1 + bot/uv.lock | 27 ++++++++++++++++++++++++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/bot/pyproject.toml b/bot/pyproject.toml index af0c38c6..d362494d 100644 --- a/bot/pyproject.toml +++ b/bot/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ "ddgs>=9.0.0", "gradio>=6.6.0", "py-machineid>=1.0.0", + "ruff>=0.15.1", ] [project.optional-dependencies] diff --git a/bot/uv.lock b/bot/uv.lock index 39c55752..34f7a46b 100644 --- a/bot/uv.lock +++ b/bot/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.11, <=3.14" resolution-markers = [ "python_full_version >= '3.14'", @@ -2524,6 +2524,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f6/f0/10642828a8dfb741e5f3fbaac830550a518a775c7fff6f04a007259b0548/py-1.11.0-py2.py3-none-any.whl", hash = "sha256:607c53218732647dff4acdfcd50cb62615cedf612e72d1724fb1a0cc6405b378", size = 98708, upload-time = "2021-11-04T17:17:00.152Z" }, ] +[[package]] +name = "py-machineid" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "winregistry", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f4/b0/c7fa6de7298a8f4e544929b97fa028304c0e11a4bc9500eff8689821bdbb/py_machineid-1.0.0.tar.gz", hash = "sha256:8a902a00fae8c6d6433f463697c21dc4ce98c6e55a2e0535c0273319acb0047a", size = 4629, upload-time = "2025-12-02T16:12:54.286Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/be/76/1ed8375cb1212824c57eb706e1f09f3f2ca4ed12b8d56b28a160e2d53505/py_machineid-1.0.0-py3-none-any.whl", hash = "sha256:910df0d5f2663bcf6739d835c4949f4e9cc6bb090a58b3dd766e12e5f768e3b9", size = 4926, upload-time = "2025-12-02T16:12:20.584Z" }, +] + [[package]] name = "pyagfs" version = "1.4.0" @@ -3648,6 +3660,7 @@ dependencies = [ { name = "msgpack" }, { name = "openviking" }, { name = "prompt-toolkit" }, + { name = "py-machineid" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "pygments" }, @@ -3655,6 +3668,7 @@ dependencies = [ { name = "python-socks" }, { name = "readability-lxml" }, { name = "rich" }, + { name = "ruff" }, { name = "socksio" }, { name = "typer" }, { name = "uvicorn" }, @@ -3741,6 +3755,7 @@ requires-dist = [ { name = "opensandbox-server", marker = "extra == 'sandbox'", specifier = ">=0.1.0" }, { name = "openviking", specifier = ">=0.1.18" }, { name = "prompt-toolkit", specifier = ">=3.0.0" }, + { name = "py-machineid", specifier = ">=1.0.0" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "pydantic-settings", specifier = ">=2.0.0" }, { name = "pygments", specifier = ">=2.16.0" }, @@ -3754,6 +3769,7 @@ requires-dist = [ { name = "qq-botpy", marker = "extra == 'qq'", specifier = ">=1.0.0" }, { name = "readability-lxml", specifier = ">=0.8.0" }, { name = "rich", specifier = ">=13.0.0" }, + { name = "ruff", specifier = ">=0.15.1" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.1.0" }, { name = "slack-sdk", marker = "extra == 'full'", specifier = ">=3.26.0" }, { name = "slack-sdk", marker = "extra == 'slack'", specifier = ">=3.26.0" }, @@ -3901,6 +3917,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e1/07/c6fe3ad3e685340704d314d765b7912993bcb8dc198f0e7a89382d37974b/win32_setctime-1.2.0-py3-none-any.whl", hash = "sha256:95d644c4e708aba81dc3704a116d8cbc974d70b3bdb8be1d150e36be6e9d1390", size = 4083, upload-time = "2024-12-07T15:28:26.465Z" }, ] +[[package]] +name = "winregistry" +version = "2.1.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/ab/6c646e9b2b6bd30c10fa69bf6733cc6eb519611552f0915c18e1522ac743/winregistry-2.1.4.tar.gz", hash = "sha256:c557e6ec26a2827625451bbcee394b2d4c1fbccb6d7dca066335f397c0be1f89", size = 9877, upload-time = "2026-02-15T09:54:06.131Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/73/de41b09e8988451f7911ee31c39ba0fc6430ba60667110006b324e96a50b/winregistry-2.1.4-py3-none-any.whl", hash = "sha256:8a2ead541bf9737d2cf4aa7c36363054a792a9f5c8fdaf0f7e3739ff9ddb092f", size = 8906, upload-time = "2026-02-15T09:54:05.355Z" }, +] + [[package]] name = "wrapt" version = "1.16.0" From fa86efc8448d05918591d90cdaa4dcbfb5557807 Mon Sep 17 00:00:00 2001 From: chenjunwen Date: Fri, 6 Mar 2026 11:02:08 +0800 Subject: [PATCH 21/21] ruff format . --- bot/tests/conftest.py | 1 - bot/tests/test_chat_commands.py | 133 ---------- bot/tests/test_chat_functionality.py | 248 ++++++++++++++++++ bot/vikingbot/agent/context.py | 5 +- bot/vikingbot/agent/loop.py | 9 +- bot/vikingbot/agent/memory.py | 5 +- bot/vikingbot/bus/events.py | 1 + bot/vikingbot/bus/queue.py | 4 +- bot/vikingbot/channels/base.py | 4 +- bot/vikingbot/channels/chat.py | 7 +- bot/vikingbot/channels/manager.py | 6 +- bot/vikingbot/channels/openapi.py | 5 +- bot/vikingbot/channels/openapi_models.py | 20 +- bot/vikingbot/channels/qq.py | 4 +- bot/vikingbot/channels/single_turn.py | 1 + bot/vikingbot/cli/commands.py | 52 ++-- bot/vikingbot/config/loader.py | 2 + bot/vikingbot/config/schema.py | 7 +- .../hooks/builtins/openviking_hooks.py | 12 +- bot/vikingbot/integrations/langfuse.py | 17 +- bot/vikingbot/openviking_mount/ov_server.py | 13 +- .../openviking_mount/session_integration.py | 1 + bot/vikingbot/providers/litellm_provider.py | 7 +- bot/vikingbot/sandbox/backends/aiosandbox.py | 4 +- bot/vikingbot/sandbox/backends/direct.py | 2 +- bot/vikingbot/sandbox/backends/opensandbox.py | 4 +- bot/vikingbot/sandbox/backends/srt.py | 4 +- bot/vikingbot/utils/tracing.py | 30 ++- openviking/async_client.py | 16 +- openviking/client/local.py | 12 +- openviking/eval/ragas/__init__.py | 8 +- openviking/eval/ragas/base.py | 12 +- openviking/eval/ragas/generator.py | 11 +- openviking/eval/ragas/pipeline.py | 4 +- openviking/eval/ragas/play_recorder.py | 15 +- openviking/eval/ragas/rag_eval.py | 27 +- openviking/eval/ragas/record_analysis.py | 8 +- openviking/eval/recorder/playback.py | 1 + openviking/models/embedder/jina_embedders.py | 1 - .../parse/parsers/code/ast/extractor.py | 29 +- .../parse/parsers/code/ast/languages/cpp.py | 33 ++- .../parse/parsers/code/ast/languages/go.py | 6 +- .../parse/parsers/code/ast/languages/java.py | 14 +- .../parse/parsers/code/ast/languages/js_ts.py | 35 ++- .../parsers/code/ast/languages/python.py | 6 +- .../parse/parsers/code/ast/languages/rust.py | 6 +- openviking/parse/parsers/code/ast/skeleton.py | 17 +- openviking/parse/parsers/code/code.py | 10 +- openviking/parse/tree_builder.py | 6 +- openviking/server/bootstrap.py | 4 +- openviking/server/routers/bot.py | 2 +- openviking/service/resource_service.py | 10 +- openviking/session/__init__.py | 7 +- .../vectordb/project/vikingdb_project.py | 14 +- .../storage/vectordb_adapters/factory.py | 1 + openviking/storage/viking_fs.py | 5 +- openviking/utils/embedding_utils.py | 71 +++-- openviking/utils/resource_processor.py | 36 +-- openviking/utils/summarizer.py | 15 +- .../utils/config/vectordb_config.py | 4 +- tests/client/test_import_export.py | 5 +- tests/integration/test_add_resource_index.py | 104 ++++---- tests/parse/test_ast_extractor.py | 90 ++++--- tests/server/test_api_filesystem.py | 16 +- tests/server/test_api_relations.py | 8 +- tests/server/test_api_resources.py | 16 +- tests/server/test_api_search.py | 4 +- tests/server/test_server_health.py | 4 +- tests/storage/mock_backend.py | 31 ++- tests/storage/test_vectordb_adaptor.py | 26 +- .../test_vectordb_collection_loading.py | 71 +++-- tests/test_code_hosting_utils.py | 4 +- tests/utils/mock_agfs.py | 48 ++-- tests/vectordb/test_openviking_vectordb.py | 8 +- 74 files changed, 870 insertions(+), 619 deletions(-) delete mode 100644 bot/tests/test_chat_commands.py create mode 100644 bot/tests/test_chat_functionality.py diff --git a/bot/tests/conftest.py b/bot/tests/conftest.py index c9cf65ca..9ace5753 100644 --- a/bot/tests/conftest.py +++ b/bot/tests/conftest.py @@ -194,4 +194,3 @@ async def client_with_resource( result = await client.add_resource(path=str(sample_markdown_file), reason="Test resource") uri = result.get("root_uri", "") yield client, uri - diff --git a/bot/tests/test_chat_commands.py b/bot/tests/test_chat_commands.py deleted file mode 100644 index a88e5690..00000000 --- a/bot/tests/test_chat_commands.py +++ /dev/null @@ -1,133 +0,0 @@ -# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 -"""Tests for vikingbot chat commands""" - -import subprocess -import sys -from pathlib import Path - -import pytest - - -def test_vikingbot_chat_help(): - """Test that vikingbot chat --help shows correct options""" - # Run vikingbot chat --help - result = subprocess.run( - [sys.executable, "-m", "vikingbot.cli.commands", "chat", "--help"], - capture_output=True, - text=True, - ) - - # Check exit code - assert result.returncode == 0, f"Command failed: {result.stderr}" - - # Check that expected options are present - assert "--message" in result.stdout or "-m" in result.stdout - assert "--session" in result.stdout or "-s" in result.stdout - assert "--markdown" in result.stdout - assert "--no-markdown" in result.stdout - assert "--logs" in result.stdout - assert "--no-logs" in result.stdout - - -def test_vikingbot_command_exists(): - """Test that vikingbot command can be invoked""" - # Just check that the main module can be imported and shows help - result = subprocess.run( - [sys.executable, "-m", "vikingbot.cli.commands", "--help"], - capture_output=True, - text=True, - ) - - assert result.returncode == 0 - assert "chat" in result.stdout - assert "Interact with the agent directly" in result.stdout - - -def test_vikingbot_gateway_help(): - """Test that gateway command help works""" - result = subprocess.run( - [sys.executable, "-m", "vikingbot.cli.commands", "gateway", "--help"], - capture_output=True, - text=True, - ) - - assert result.returncode == 0 - assert "--port" in result.stdout or "-p" in result.stdout - - -def test_single_turn_channel_import(): - """Test that SingleTurnChannel can be imported""" - from vikingbot.channels.single_turn import SingleTurnChannel, SingleTurnChannelConfig - - assert SingleTurnChannel is not None - assert SingleTurnChannelConfig is not None - - -def test_chat_channel_import(): - """Test that ChatChannel can be imported""" - from vikingbot.channels.chat import ChatChannel, ChatChannelConfig - - assert ChatChannel is not None - assert ChatChannelConfig is not None - - -def test_prepare_agent_channel_function(): - """Test that prepare_agent_channel function exists and can be imported""" - from vikingbot.cli.commands import prepare_agent_channel - - assert prepare_agent_channel is not None - assert callable(prepare_agent_channel) - - -def test_chat_command_function_exists(): - """Test that the chat command function is registered""" - from vikingbot.cli.commands import app - - # Check that we can get the chat command info - # Typer stores commands differently, let's just verify the chat attribute exists - assert hasattr(app, "commands") or hasattr(app, "registered_commands") - - # Try calling the chat command with --help as a better verification - import subprocess - import sys - - result = subprocess.run( - [sys.executable, "-m", "vikingbot.cli.commands", "chat", "--help"], - capture_output=True, - text=True, - ) - - assert result.returncode == 0 - assert "Interact with the agent directly" in result.stdout - - -@pytest.mark.integration -def test_chat_single_turn_dry_run(): - """ - Dry run test for single-turn chat (doesn't actually call LLM) - This tests the infrastructure without requiring API keys - """ - # This test just verifies the modules load correctly - # A full integration test would need API keys and a running agent - from vikingbot.channels.single_turn import SingleTurnChannel, SingleTurnChannelConfig - from vikingbot.bus.queue import MessageBus - - config = SingleTurnChannelConfig() - bus = MessageBus() - - # Just verify we can instantiate the channel without errors - channel = SingleTurnChannel( - config, - bus, - workspace_path=Path("/tmp"), - message="Hello, test", - session_id="test-session", - markdown=True, - ) - - assert channel is not None - assert channel.name == "single_turn" - assert channel.message == "Hello, test" - assert channel.session_id == "test-session" - diff --git a/bot/tests/test_chat_functionality.py b/bot/tests/test_chat_functionality.py new file mode 100644 index 00000000..f163fe19 --- /dev/null +++ b/bot/tests/test_chat_functionality.py @@ -0,0 +1,248 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for vikingbot chat functionality - single message and interactive modes.""" + +import tempfile +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from vikingbot.bus.events import OutboundMessage +from vikingbot.bus.queue import MessageBus +from vikingbot.channels.chat import ChatChannel, ChatChannelConfig +from vikingbot.channels.single_turn import SingleTurnChannel, SingleTurnChannelConfig +from vikingbot.cli.commands import prepare_agent_channel +from vikingbot.config.schema import SessionKey + + +@pytest.fixture +def temp_workspace(): + """Create a temporary workspace directory.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + + +@pytest.fixture +def message_bus(): + """Create a MessageBus instance.""" + return MessageBus() + + +@pytest.fixture +def config(temp_workspace): + """Create a mock config.""" + config = MagicMock() + config.workspace_path = temp_workspace + config.bot_data_path = temp_workspace / "bot_data" + config.sandbox = MagicMock() + config.sandbox.backend = "direct" + config.sandbox.mode = "user" + config.agents = MagicMock() + config.agents.model = "test-model" + config.agents.api_key = "test-key" + config.agents.api_base = None + config.agents.provider = "test-provider" + config.agents.max_tool_iterations = 10 + config.agents.memory_window = 50 + config.agents.gen_image_model = "test-image-model" + config.tools = MagicMock() + config.tools.web = MagicMock() + config.tools.web.search = MagicMock() + config.tools.web.search.api_key = None + config.tools.exec = MagicMock() + config.hooks = [] + config.heartbeat = MagicMock() + config.heartbeat.enabled = False + config.heartbeat.interval_seconds = 60 + config.langfuse = MagicMock() + config.langfuse.enabled = False + config.providers = MagicMock() + return config + + +class TestSingleTurnChannel: + """Tests for SingleTurnChannel (vikingbot chat -m xxx).""" + + def test_single_turn_channel_initialization(self, message_bus, temp_workspace): + """Test that SingleTurnChannel can be initialized correctly.""" + config = SingleTurnChannelConfig() + channel = SingleTurnChannel( + config, + message_bus, + workspace_path=temp_workspace, + message="Hello, test", + session_id="test-session", + markdown=True, + ) + + assert channel is not None + assert channel.name == "single_turn" + assert channel.message == "Hello, test" + assert channel.session_id == "test-session" + + @pytest.mark.asyncio + async def test_single_turn_channel_receives_response(self, message_bus, temp_workspace): + """Test that SingleTurnChannel can receive and store responses.""" + config = SingleTurnChannelConfig() + test_message = "Hello, test" + channel = SingleTurnChannel( + config, + message_bus, + workspace_path=temp_workspace, + message=test_message, + session_id="test-session", + markdown=True, + ) + + # Create a test response + session_key = SessionKey(type="cli", channel_id="default", chat_id="test-session") + test_response = "This is a test response from the bot" + + # Send the response + await channel.send( + OutboundMessage( + session_key=session_key, + content=test_response, + ) + ) + + # Check that the response was stored + assert channel._last_response == test_response + assert channel._response_received.is_set() + + +class TestChatChannel: + """Tests for ChatChannel (interactive vikingbot chat).""" + + def test_chat_channel_initialization(self, message_bus, temp_workspace): + """Test that ChatChannel can be initialized correctly.""" + config = ChatChannelConfig() + channel = ChatChannel( + config, + message_bus, + workspace_path=temp_workspace, + session_id="test-session", + markdown=True, + logs=False, + ) + + assert channel is not None + assert channel.name == "chat" + assert channel.session_id == "test-session" + + @pytest.mark.asyncio + async def test_chat_channel_send_response(self, message_bus, temp_workspace): + """Test that ChatChannel can receive and store responses.""" + config = ChatChannelConfig() + channel = ChatChannel( + config, + message_bus, + workspace_path=temp_workspace, + session_id="test-session", + markdown=True, + logs=False, + ) + + # Start the channel in background (it will wait for input) + channel._running = True + + # Create a test response + session_key = SessionKey(type="cli", channel_id="default", chat_id="test-session") + test_response = "This is a test response from the bot" + + # Send the response + await channel.send( + OutboundMessage( + session_key=session_key, + content=test_response, + ) + ) + + # Check that the response was stored + assert channel._last_response == test_response + assert channel._response_received.is_set() + + +class TestPrepareAgentChannel: + """Tests for prepare_agent_channel function.""" + + def test_prepare_agent_channel_single_message(self, message_bus, config, temp_workspace): + """Test prepare_agent_channel with a single message (vikingbot chat -m xxx).""" + test_message = "Hello, this is a single message" + session_id = "test-session-123" + + channels = prepare_agent_channel( + config, + message_bus, + message=test_message, + session_id=session_id, + markdown=True, + logs=False, + ) + + assert channels is not None + # Check that we have a SingleTurnChannel + assert len(channels.channels) == 1 + # channels is a dict, get the first value + channel = next(iter(channels.channels.values())) + assert channel.name == "single_turn" + assert channel.message == test_message + assert channel.session_id == session_id + + def test_prepare_agent_channel_interactive(self, message_bus, config, temp_workspace): + """Test prepare_agent_channel for interactive mode (vikingbot chat).""" + session_id = "test-session-456" + + channels = prepare_agent_channel( + config, + message_bus, + message=None, # None means interactive mode + session_id=session_id, + markdown=True, + logs=True, + ) + + assert channels is not None + # Check that we have a ChatChannel + assert len(channels.channels) == 1 + # channels is a dict, get the first value + channel = next(iter(channels.channels.values())) + assert channel.name == "chat" + assert channel.session_id == session_id + assert channel.logs is True + + +class TestChatCommandCLI: + """Tests for the chat command CLI interface.""" + + def test_vikingbot_chat_help(self): + """Test that vikingbot chat --help shows correct options.""" + import subprocess + import sys + + result = subprocess.run( + [sys.executable, "-m", "vikingbot.cli.commands", "chat", "--help"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert "--message" in result.stdout or "-m" in result.stdout + assert "--session" in result.stdout or "-s" in result.stdout + assert "--markdown" in result.stdout + assert "--no-markdown" in result.stdout + + def test_vikingbot_gateway_help(self): + """Test that gateway command help works.""" + import subprocess + import sys + + result = subprocess.run( + [sys.executable, "-m", "vikingbot.cli.commands", "gateway", "--help"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert "--port" in result.stdout or "-p" in result.stdout diff --git a/bot/vikingbot/agent/context.py b/bot/vikingbot/agent/context.py index 0251a9ab..23684b96 100644 --- a/bot/vikingbot/agent/context.py +++ b/bot/vikingbot/agent/context.py @@ -93,14 +93,15 @@ async def build_system_prompt( ) # Add session context - session_context ="## Current Session" + session_context = "## Current Session" if session_key and session_key.type: session_context += f"\nChannel: {session_key.type}" if self._is_group_chat: session_context += ( f"\n**Group chat session.** Current user ID: {self._sender_id}\n" f"Multiple users can participate in this conversation. Each user message is prefixed with the user ID in brackets like @. " - f"You should pay attention to who is speaking to understand the context. ") + f"You should pay attention to who is speaking to understand the context. " + ) parts.append(session_context) # Viking user profile diff --git a/bot/vikingbot/agent/loop.py b/bot/vikingbot/agent/loop.py index 8a71adf8..bd1427ff 100644 --- a/bot/vikingbot/agent/loop.py +++ b/bot/vikingbot/agent/loop.py @@ -249,7 +249,7 @@ async def _run_agent_loop( model=self.model, session_id=session_key.safe_name(), ) - if response.usage: + if response.usage: cur_token = response.usage self._token_usage["prompt_tokens"] += cur_token["prompt_tokens"] self._token_usage["completion_tokens"] += cur_token["completion_tokens"] @@ -407,8 +407,13 @@ async def _process_message(self, msg: InboundMessage) -> OutboundMessage | None: message_workspace = self.workspace from vikingbot.agent.context import ContextBuilder + message_context = ContextBuilder( - message_workspace, sandbox_manager=self.sandbox_manager, sender_id=msg.sender_id, is_group_chat=is_group_chat, eval=self._eval + message_workspace, + sandbox_manager=self.sandbox_manager, + sender_id=msg.sender_id, + is_group_chat=is_group_chat, + eval=self._eval, ) # Build initial messages (use get_history for LLM-formatted messages) diff --git a/bot/vikingbot/agent/memory.py b/bot/vikingbot/agent/memory.py index 67ec7c71..547f7c50 100644 --- a/bot/vikingbot/agent/memory.py +++ b/bot/vikingbot/agent/memory.py @@ -53,10 +53,7 @@ async def get_viking_memory_context(self, current_message: str, workspace_id: st return "" user_memory = self._parse_viking_memory(result["user_memory"]) agent_memory = self._parse_viking_memory(result["agent_memory"]) - return ( - f"### user memories:\n{user_memory}\n" - f"### agent memories:\n{agent_memory}" - ) + return f"### user memories:\n{user_memory}\n### agent memories:\n{agent_memory}" async def get_viking_user_profile(self, workspace_id: str, user_id: str) -> str: client = await VikingClient.create(agent_id=workspace_id) diff --git a/bot/vikingbot/bus/events.py b/bot/vikingbot/bus/events.py index c1e8fca6..147c0eab 100644 --- a/bot/vikingbot/bus/events.py +++ b/bot/vikingbot/bus/events.py @@ -10,6 +10,7 @@ class OutboundEventType(str, Enum): """Type of outbound message/event.""" + RESPONSE = "response" # Normal response message TOOL_CALL = "tool_call" # Tool being called TOOL_RESULT = "tool_result" # Result from tool execution diff --git a/bot/vikingbot/bus/queue.py b/bot/vikingbot/bus/queue.py index 83c96e97..1e4067a6 100644 --- a/bot/vikingbot/bus/queue.py +++ b/bot/vikingbot/bus/queue.py @@ -26,7 +26,7 @@ def __init__(self): async def publish_inbound(self, msg: InboundMessage) -> None: """Publish a message from a channel to the agent.""" - #print(f'publish_inbound={msg}') + # print(f'publish_inbound={msg}') await self.inbound.put(msg) async def consume_inbound(self) -> InboundMessage: @@ -35,7 +35,7 @@ async def consume_inbound(self) -> InboundMessage: async def publish_outbound(self, msg: OutboundMessage) -> None: """Publish a response from the agent to channels.""" - #print(f'publish_outbound={msg}') + # print(f'publish_outbound={msg}') await self.outbound.put(msg) async def consume_outbound(self) -> OutboundMessage: diff --git a/bot/vikingbot/channels/base.py b/bot/vikingbot/channels/base.py index abf266d8..98f12fd4 100644 --- a/bot/vikingbot/channels/base.py +++ b/bot/vikingbot/channels/base.py @@ -138,7 +138,9 @@ async def _handle_message( msg = InboundMessage( session_key=SessionKey( - type=str(getattr(self.channel_type, 'value', self.channel_type)), channel_id=self.channel_id, chat_id=chat_id + type=str(getattr(self.channel_type, "value", self.channel_type)), + channel_id=self.channel_id, + chat_id=chat_id, ), sender_id=str(sender_id), content=content, diff --git a/bot/vikingbot/channels/chat.py b/bot/vikingbot/channels/chat.py index 68408884..03ff0f70 100644 --- a/bot/vikingbot/channels/chat.py +++ b/bot/vikingbot/channels/chat.py @@ -82,7 +82,12 @@ async def send(self, msg: OutboundMessage) -> None: console.print("[bold red]Bot:[/bold red]") from rich.markdown import Markdown from rich.text import Text - body = Markdown(content, style="red") if self.markdown else Text(content, style=Style(color="red")) + + body = ( + Markdown(content, style="red") + if self.markdown + else Text(content, style=Style(color="red")) + ) console.print(body) console.print() diff --git a/bot/vikingbot/channels/manager.py b/bot/vikingbot/channels/manager.py index 9da5892a..673d78f3 100644 --- a/bot/vikingbot/channels/manager.py +++ b/bot/vikingbot/channels/manager.py @@ -146,7 +146,7 @@ def add_channel_from_config( logger.warning( f"Channel {channel_config.type} not available: {e}. " f"Install with: uv pip install 'vikingbot[{channel_type}]' " - f"(or uv pip install -e \".[{channel_type}]\" for local dev)" + f'(or uv pip install -e ".[{channel_type}]" for local dev)' ) def load_channels_from_config( @@ -162,7 +162,9 @@ def load_channels_from_config( self.add_channel_from_config( channel_config, workspace_path=workspace_path, - groq_api_key=config.providers.groq.api_key if hasattr(config.providers, "groq") else None, + groq_api_key=config.providers.groq.api_key + if hasattr(config.providers, "groq") + else None, ) async def _start_channel(self, name: str, channel: BaseChannel) -> None: diff --git a/bot/vikingbot/channels/openapi.py b/bot/vikingbot/channels/openapi.py index 74f8ee1e..488d8db6 100644 --- a/bot/vikingbot/channels/openapi.py +++ b/bot/vikingbot/channels/openapi.py @@ -389,10 +389,7 @@ async def event_generator(): except Exception as e: logger.exception(f"Error in stream generator: {e}") - error_event = ChatStreamEvent( - event=EventType.RESPONSE, - data={"error": str(e)} - ) + error_event = ChatStreamEvent(event=EventType.RESPONSE, data={"error": str(e)}) yield f"data: {error_event.model_dump_json()}\n\n" finally: self._pending.pop(session_id, None) diff --git a/bot/vikingbot/channels/openapi_models.py b/bot/vikingbot/channels/openapi_models.py index c4d36bb1..eeaf3b53 100644 --- a/bot/vikingbot/channels/openapi_models.py +++ b/bot/vikingbot/channels/openapi_models.py @@ -31,17 +31,23 @@ class ChatMessage(BaseModel): role: MessageRole = Field(..., description="Role of the message sender") content: str = Field(..., description="Message content") - timestamp: Optional[datetime] = Field(default_factory=datetime.now, description="Message timestamp") + timestamp: Optional[datetime] = Field( + default_factory=datetime.now, description="Message timestamp" + ) class ChatRequest(BaseModel): """Request body for chat endpoint.""" message: str = Field(..., description="User message to send", min_length=1) - session_id: Optional[str] = Field(default="default", description="Session ID (optional, will create new if not provided)") + session_id: Optional[str] = Field( + default="default", description="Session ID (optional, will create new if not provided)" + ) user_id: Optional[str] = Field(default=None, description="User identifier (optional)") stream: bool = Field(default=False, description="Whether to stream the response") - context: Optional[List[ChatMessage]] = Field(default=None, description="Additional context messages") + context: Optional[List[ChatMessage]] = Field( + default=None, description="Additional context messages" + ) class ChatResponse(BaseModel): @@ -49,7 +55,9 @@ class ChatResponse(BaseModel): session_id: str = Field(..., description="Session ID") message: str = Field(..., description="Assistant's response message") - events: Optional[List[Dict[str, Any]]] = Field(default=None, description="Intermediate events (thinking, tool calls)") + events: Optional[List[Dict[str, Any]]] = Field( + default=None, description="Intermediate events (thinking, tool calls)" + ) timestamp: datetime = Field(default_factory=datetime.now, description="Response timestamp") @@ -74,7 +82,9 @@ class SessionCreateRequest(BaseModel): """Request to create a new session.""" user_id: Optional[str] = Field(default=None, description="User identifier") - metadata: Optional[Dict[str, Any]] = Field(default=None, description="Optional session metadata") + metadata: Optional[Dict[str, Any]] = Field( + default=None, description="Optional session metadata" + ) class SessionCreateResponse(BaseModel): diff --git a/bot/vikingbot/channels/qq.py b/bot/vikingbot/channels/qq.py index 49eda8fd..49aed93a 100644 --- a/bot/vikingbot/channels/qq.py +++ b/bot/vikingbot/channels/qq.py @@ -60,7 +60,9 @@ def __init__(self, config: QQChannelConfig, bus: MessageBus, **kwargs): async def start(self) -> None: """Start the QQ bot.""" if not QQ_AVAILABLE: - logger.exception("QQ SDK not installed. Install with: uv pip install 'vikingbot[qq]' (or uv pip install -e \".[qq]\" for local dev)") + logger.exception( + "QQ SDK not installed. Install with: uv pip install 'vikingbot[qq]' (or uv pip install -e \".[qq]\" for local dev)" + ) return if not self.config.app_id or not self.config.secret: diff --git a/bot/vikingbot/channels/single_turn.py b/bot/vikingbot/channels/single_turn.py index 4294e8e8..5c5dc326 100644 --- a/bot/vikingbot/channels/single_turn.py +++ b/bot/vikingbot/channels/single_turn.py @@ -77,6 +77,7 @@ async def start(self) -> None: from vikingbot.cli.commands import console from rich.markdown import Markdown from rich.text import Text + content = self._last_response or "" body = Markdown(content) if self.markdown else Text(content) console.print(body) diff --git a/bot/vikingbot/cli/commands.py b/bot/vikingbot/cli/commands.py index 98c2756e..9256fb4b 100644 --- a/bot/vikingbot/cli/commands.py +++ b/bot/vikingbot/cli/commands.py @@ -36,7 +36,12 @@ # Create sandbox manager from vikingbot.sandbox.manager import SandboxManager from vikingbot.session.manager import SessionManager -from vikingbot.utils.helpers import get_source_workspace_path, set_bot_data_path, get_history_path, get_bridge_path +from vikingbot.utils.helpers import ( + get_source_workspace_path, + set_bot_data_path, + get_history_path, + get_bridge_path, +) app = typer.Typer( name="vikingbot", @@ -55,6 +60,7 @@ def get_or_create_machine_id() -> str: """ try: from machineid import machine_id + return machine_id() except ImportError: # Fallback if py-machineid is not installed @@ -66,7 +72,6 @@ def get_or_create_machine_id() -> str: return "default" - def _init_bot_data(config): """Initialize bot data directory and set global paths.""" set_bot_data_path(config.bot_data_path) @@ -190,11 +195,10 @@ def main( pass -def _make_provider(config, langfuse_client: None = None): +def _make_provider(config, langfuse_client: None = None): """Create LiteLLMProvider from config. Allows starting without API key.""" from vikingbot.providers.litellm_provider import LiteLLMProvider - config = load_config() p = config.agents @@ -248,6 +252,7 @@ def gateway( # Create FastAPI app for OpenAPI from fastapi import FastAPI + fastapi_app = FastAPI( title="Vikingbot OpenAPI", description="HTTP API for Vikingbot chat", @@ -255,7 +260,9 @@ def gateway( ) cron = prepare_cron(bus) - channels = prepare_channel(config, bus, fastapi_app=fastapi_app, enable_openapi=True, openapi_port=port) + channels = prepare_channel( + config, bus, fastapi_app=fastapi_app, enable_openapi=True, openapi_port=port + ) agent_loop = prepare_agent_loop(config, bus, session_manager, cron) heartbeat = prepare_heartbeat(config, agent_loop, session_manager) @@ -292,7 +299,9 @@ def prepare_agent_loop(config, bus, session_manager, cron, quiet: bool = False, if config.sandbox.backend == "direct": logger.warning("[SANDBOX] disabled (using DIRECT mode - commands run directly on host)") else: - logger.info(f"Sandbox: enabled (backend={config.sandbox.backend}, mode={config.sandbox.mode})") + logger.info( + f"Sandbox: enabled (backend={config.sandbox.backend}, mode={config.sandbox.mode})" + ) # Initialize Langfuse if enabled langfuse_client = None @@ -328,11 +337,11 @@ def prepare_agent_loop(config, bus, session_manager, cron, quiet: bool = False, session_manager=session_manager, sandbox_manager=sandbox_manager, config=config, - eval=eval + eval=eval, ) # Set the agent reference in cron if it uses the holder pattern - if hasattr(cron, '_agent_holder'): - cron._agent_holder['agent'] = agent + if hasattr(cron, "_agent_holder"): + cron._agent_holder["agent"] = agent return agent @@ -393,7 +402,9 @@ async def on_cron_job(job: CronJob) -> str | None: return cron -def prepare_channel(config, bus, fastapi_app=None, enable_openapi: bool = False, openapi_port: int = 18790): +def prepare_channel( + config, bus, fastapi_app=None, enable_openapi: bool = False, openapi_port: int = 18790 +): """Prepare channels for the bot. Args: @@ -485,11 +496,20 @@ def _thinking_ctx(logs: bool): """Return a context manager for showing thinking spinner.""" if logs: from contextlib import nullcontext + return nullcontext() return console.status("[dim]vikingbot is thinking...[/dim]", spinner="dots") -def prepare_agent_channel(config, bus, message: str | None, session_id: str, markdown: bool, logs: bool, eval: bool = False): +def prepare_agent_channel( + config, + bus, + message: str | None, + session_id: str, + markdown: bool, + logs: bool, + eval: bool = False, +): """Prepare channel for agent command.""" from vikingbot.channels.chat import ChatChannel, ChatChannelConfig from vikingbot.channels.single_turn import SingleTurnChannel, SingleTurnChannelConfig @@ -559,7 +579,9 @@ def chat( session_id = get_or_create_machine_id() cron = prepare_cron(bus, quiet=is_single_turn) channels = prepare_agent_channel(config, bus, message, session_id, markdown, logs, eval) - agent_loop = prepare_agent_loop(config, bus, session_manager, cron, quiet=is_single_turn, eval=eval) + agent_loop = prepare_agent_loop( + config, bus, session_manager, cron, quiet=is_single_turn, eval=eval + ) async def run(): if is_single_turn: @@ -569,10 +591,7 @@ async def run(): task_agent = asyncio.create_task(agent_loop.run()) # Wait for channels to complete (it will complete after getting response) - done, pending = await asyncio.wait( - [task_channels], - return_when=asyncio.FIRST_COMPLETED - ) + done, pending = await asyncio.wait([task_channels], return_when=asyncio.FIRST_COMPLETED) # Cancel all other tasks for task in pending: @@ -949,6 +968,7 @@ def status(): try: from vikingbot.cli.test_commands import test_app + app.add_typer(test_app, name="test") except ImportError: # If test commands not available, don't add them diff --git a/bot/vikingbot/config/loader.py b/bot/vikingbot/config/loader.py index 1673136f..d32b7438 100644 --- a/bot/vikingbot/config/loader.py +++ b/bot/vikingbot/config/loader.py @@ -129,6 +129,7 @@ def _merge_vlm_model_config(bot_data: dict, vlm_data: dict) -> None: bot_data["agents"]["api_base"] = vlm_data.get("api_base", "") bot_data["agents"]["api_key"] = vlm_data.get("api_key", "") + def _merge_ov_server_config(bot_data: dict, ov_data: dict) -> None: """ Merge ov_server config into bot config. @@ -144,6 +145,7 @@ def _merge_ov_server_config(bot_data: dict, ov_data: dict) -> None: else: bot_data["mode"] = "local" + def save_config(config: Config, config_path: Path | None = None) -> None: """ Save configuration to ov.conf's bot field, preserving other sections. diff --git a/bot/vikingbot/config/schema.py b/bot/vikingbot/config/schema.py index ef07f802..653f76f4 100644 --- a/bot/vikingbot/config/schema.py +++ b/bot/vikingbot/config/schema.py @@ -100,7 +100,7 @@ class FeishuChannelConfig(BaseChannelConfig): app_secret: str = "" encrypt_key: str = "" verification_token: str = "" - allow_from: list[str] = Field(default_factory=list) ## 允许更新Agent对话的Feishu用户ID列表 + allow_from: list[str] = Field(default_factory=list) ## 允许更新Agent对话的Feishu用户ID列表 def channel_id(self) -> str: # Use app_id directly as the ID @@ -582,7 +582,9 @@ class Config(BaseSettings): agents: AgentsConfig = Field(default_factory=AgentsConfig) channels: list[Any] = Field(default_factory=list) - providers: ProvidersConfig = Field(default_factory=ProvidersConfig, deprecated=True) # Deprecated: Use ov.conf vlm config instead + providers: ProvidersConfig = Field( + default_factory=ProvidersConfig, deprecated=True + ) # Deprecated: Use ov.conf vlm config instead gateway: GatewayConfig = Field(default_factory=GatewayConfig) tools: ToolsConfig = Field(default_factory=ToolsConfig) ov_server: OpenVikingConfig = Field(default_factory=OpenVikingConfig) @@ -629,6 +631,7 @@ def ov_data_path(self) -> Path: def _get_vlm_config(self) -> Optional[Dict[str, Any]]: """Get vlm config from OpenVikingConfig. Returns (vlm_config_dict).""" from openviking_cli.utils.config import get_openviking_config + ov_config = get_openviking_config() if hasattr(ov_config, "vlm"): diff --git a/bot/vikingbot/hooks/builtins/openviking_hooks.py b/bot/vikingbot/hooks/builtins/openviking_hooks.py index 3e3daded..b06e1c82 100644 --- a/bot/vikingbot/hooks/builtins/openviking_hooks.py +++ b/bot/vikingbot/hooks/builtins/openviking_hooks.py @@ -58,14 +58,20 @@ async def execute(self, context: HookContext, **kwargs) -> Any: try: allow_from = self._get_channel_allow_from(session_id) - filtered_messages = self._filter_messages_by_sender(vikingbot_session.messages, allow_from) + filtered_messages = self._filter_messages_by_sender( + vikingbot_session.messages, allow_from + ) if not filtered_messages: - logger.info(f"No messages to commit openviking for session {session_id} (allow_from filter applied)") + logger.info( + f"No messages to commit openviking for session {session_id} (allow_from filter applied)" + ) return {"success": True, "message": "No messages matched allow_from filter"} client = await self._get_client(context.workspace_id) - result = await client.commit(session_id, filtered_messages, load_config().ov_server.admin_user_id) + result = await client.commit( + session_id, filtered_messages, load_config().ov_server.admin_user_id + ) return result except Exception as e: logger.exception(f"Failed to add message to OpenViking: {e}") diff --git a/bot/vikingbot/integrations/langfuse.py b/bot/vikingbot/integrations/langfuse.py index c6ea7ad2..f071ea3d 100644 --- a/bot/vikingbot/integrations/langfuse.py +++ b/bot/vikingbot/integrations/langfuse.py @@ -12,6 +12,7 @@ try: from langfuse import Langfuse from langfuse import propagate_attributes as _propagate_attributes + propagate_attributes = _propagate_attributes except ImportError: pass @@ -36,12 +37,16 @@ def __init__( return if Langfuse is None: - logger.warning("Langfuse not installed. Install with: uv pip install vikingbot[langfuse] (or uv pip install -e \".[langfuse]\" for local dev). Configure in ~/.openviking/ov.conf under bot.langfuse") + logger.warning( + 'Langfuse not installed. Install with: uv pip install vikingbot[langfuse] (or uv pip install -e ".[langfuse]" for local dev). Configure in ~/.openviking/ov.conf under bot.langfuse' + ) self.enabled = False return if not secret_key: - logger.warning("Langfuse enabled but no secret_key provided. Configure in ~/.openviking/ov.conf under bot.langfuse") + logger.warning( + "Langfuse enabled but no secret_key provided. Configure in ~/.openviking/ov.conf under bot.langfuse" + ) self.enabled = False return @@ -93,7 +98,9 @@ def propagate_attributes( yield return if not self._client: - logger.warning(f"[LANGFUSE] propagate_attributes skipped: Langfuse client not initialized") + logger.warning( + f"[LANGFUSE] propagate_attributes skipped: Langfuse client not initialized" + ) yield return @@ -115,7 +122,9 @@ def propagate_attributes( with propagate_attributes(**propagate_kwargs): yield else: - logger.warning(f"[LANGFUSE] propagate_attributes not available (SDK version may not support it)") + logger.warning( + f"[LANGFUSE] propagate_attributes not available (SDK version may not support it)" + ) yield except Exception as e: logger.debug(f"[LANGFUSE] propagate_attributes error: {e}") diff --git a/bot/vikingbot/openviking_mount/ov_server.py b/bot/vikingbot/openviking_mount/ov_server.py index 987f3df7..287af820 100644 --- a/bot/vikingbot/openviking_mount/ov_server.py +++ b/bot/vikingbot/openviking_mount/ov_server.py @@ -19,9 +19,7 @@ def __init__(self, agent_id: Optional[str] = None): self.openviking_config = openviking_config self.ov_path = config.ov_data_path if openviking_config.mode == "local": - self.client = ov.AsyncHTTPClient( - url=openviking_config.server_url - ) + self.client = ov.AsyncHTTPClient(url=openviking_config.server_url) self.agent_id = "default" self.account_id = "default" self.admin_user_id = "default" @@ -348,7 +346,12 @@ async def commit(self, session_id: str, messages: list[dict[str, Any]], user_id: # For remote mode, try to get user's API key and create a dedicated client client = self.client start = time.time() - if self.mode == "remote" and user_id and user_id != self.admin_user_id and self._apikey_manager: + if ( + self.mode == "remote" + and user_id + and user_id != self.admin_user_id + and self._apikey_manager + ): user_api_key = await self._get_or_create_user_apikey(user_id) if user_api_key: # Create a new HTTP client with user's API key @@ -360,7 +363,7 @@ async def commit(self, session_id: str, messages: list[dict[str, Any]], user_id: await client.initialize() create_res = await client.create_session() - session_id = create_res['session_id'] + session_id = create_res["session_id"] session = client.session(session_id) for message in messages: diff --git a/bot/vikingbot/openviking_mount/session_integration.py b/bot/vikingbot/openviking_mount/session_integration.py index b16b3f4d..a6ec77a5 100644 --- a/bot/vikingbot/openviking_mount/session_integration.py +++ b/bot/vikingbot/openviking_mount/session_integration.py @@ -16,6 +16,7 @@ from loguru import logger from vikingbot.utils.helpers import get_workspace_path + # 相对导入同一包内的模块 from .mount import OpenVikingMount, MountConfig, MountScope from .viking_fuse import mount_fuse, FUSEMountManager, FUSE_AVAILABLE diff --git a/bot/vikingbot/providers/litellm_provider.py b/bot/vikingbot/providers/litellm_provider.py index 84eea405..646a6779 100644 --- a/bot/vikingbot/providers/litellm_provider.py +++ b/bot/vikingbot/providers/litellm_provider.py @@ -195,10 +195,9 @@ async def chat( # Add cache read tokens if available (OpenAI/Anthropic prompt caching) # Try multiple possible field names for cached tokens - cache_read_tokens = ( - llm_response.usage.get("cache_read_input_tokens") or - llm_response.usage.get("prompt_tokens_details", {}).get("cached_tokens") - ) + cache_read_tokens = llm_response.usage.get( + "cache_read_input_tokens" + ) or llm_response.usage.get("prompt_tokens_details", {}).get("cached_tokens") if cache_read_tokens: usage_details["cache_read_input_tokens"] = cache_read_tokens diff --git a/bot/vikingbot/sandbox/backends/aiosandbox.py b/bot/vikingbot/sandbox/backends/aiosandbox.py index e15cf03c..1447ee39 100644 --- a/bot/vikingbot/sandbox/backends/aiosandbox.py +++ b/bot/vikingbot/sandbox/backends/aiosandbox.py @@ -35,7 +35,9 @@ async def start(self) -> None: self._client = AsyncSandbox(base_url=self._base_url) logger.info("[AioSandbox] Connected successfully") except ImportError: - logger.error("agent-sandbox SDK not installed. Install with: uv pip install 'vikingbot[sandbox]' (or uv pip install -e \".[sandbox]\" for local dev)") + logger.error( + "agent-sandbox SDK not installed. Install with: uv pip install 'vikingbot[sandbox]' (or uv pip install -e \".[sandbox]\" for local dev)" + ) raise except Exception as e: logger.error("[AioSandbox] Failed to start: {}", e) diff --git a/bot/vikingbot/sandbox/backends/direct.py b/bot/vikingbot/sandbox/backends/direct.py index 37859a21..a273d21f 100644 --- a/bot/vikingbot/sandbox/backends/direct.py +++ b/bot/vikingbot/sandbox/backends/direct.py @@ -29,7 +29,7 @@ async def start(self) -> None: """Start the backend (no-op for direct backend).""" self._workspace.mkdir(parents=True, exist_ok=True) self._running = True - #logger.info("Direct backend started") + # logger.info("Direct backend started") async def execute(self, command: str, timeout: int = 60, **kwargs: Any) -> str: """Execute a command directly on the host.""" diff --git a/bot/vikingbot/sandbox/backends/opensandbox.py b/bot/vikingbot/sandbox/backends/opensandbox.py index 66c1027f..712daeb7 100644 --- a/bot/vikingbot/sandbox/backends/opensandbox.py +++ b/bot/vikingbot/sandbox/backends/opensandbox.py @@ -205,7 +205,9 @@ async def start(self) -> None: logger.info("OpenSandbox created successfully") except ImportError: - logger.error("opensandbox SDK not installed. Install with: uv pip install 'vikingbot[sandbox]' (or uv pip install -e \".[sandbox]\" for local dev)") + logger.error( + "opensandbox SDK not installed. Install with: uv pip install 'vikingbot[sandbox]' (or uv pip install -e \".[sandbox]\" for local dev)" + ) raise except Exception as e: logger.error("Failed to create OpenSandbox: {}", e) diff --git a/bot/vikingbot/sandbox/backends/srt.py b/bot/vikingbot/sandbox/backends/srt.py index 22e7a393..4f09f608 100644 --- a/bot/vikingbot/sandbox/backends/srt.py +++ b/bot/vikingbot/sandbox/backends/srt.py @@ -43,9 +43,7 @@ def _generate_settings(self) -> Path: # Place settings file in workspace/sandboxes/ directory settings_path = ( - self._workspace - / "sandboxes" - / f"{self.session_key.safe_name()}-srt-settings.json" + self._workspace / "sandboxes" / f"{self.session_key.safe_name()}-srt-settings.json" ) settings_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/bot/vikingbot/utils/tracing.py b/bot/vikingbot/utils/tracing.py index 58039a08..948c30e9 100644 --- a/bot/vikingbot/utils/tracing.py +++ b/bot/vikingbot/utils/tracing.py @@ -30,7 +30,6 @@ def get_current_session_id() -> str | None: return _session_id.get() - @contextmanager def set_session_id(session_id: str | None) -> Generator[None, None, None]: """ @@ -85,6 +84,7 @@ async def process_message(msg: InboundMessage) -> Response: # session_id and user_id are automatically propagated return await handle(msg) """ + def decorator(func: Callable[..., T]) -> Callable[..., T]: span_name = name or func.__name__ @@ -101,11 +101,15 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> T: try: # Inspect the extractor's signature to determine how to call it import inspect + sig = inspect.signature(extract_session_id) - param_count = len([ - p for p in sig.parameters.values() - if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) - ]) + param_count = len( + [ + p + for p in sig.parameters.values() + if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) + ] + ) if param_count == 1 and len(args) >= 1: # Extractor expects single arg (e.g., lambda msg: ...) @@ -122,11 +126,15 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> T: if extract_user_id: try: import inspect + sig = inspect.signature(extract_user_id) - param_count = len([ - p for p in sig.parameters.values() - if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) - ]) + param_count = len( + [ + p + for p in sig.parameters.values() + if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) + ] + ) if param_count == 1 and len(args) >= 1: user_id = extract_user_id(args[-1]) @@ -140,11 +148,11 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> T: session_id = get_current_session_id() logger.debug(f"[TRACE] No session_id extracted, using context: {session_id}") else: - #logger.info(f"[TRACE] Extracted session_id: {session_id}") + # logger.info(f"[TRACE] Extracted session_id: {session_id}") pass if user_id: - #logger.info(f"[TRACE] Extracted user_id: {user_id}") + # logger.info(f"[TRACE] Extracted user_id: {user_id}") pass # Use context manager to set session_id for nested operations diff --git a/openviking/async_client.py b/openviking/async_client.py index 65fc673d..5793e172 100644 --- a/openviking/async_client.py +++ b/openviking/async_client.py @@ -213,28 +213,20 @@ async def wait_processed(self, timeout: float = None) -> Dict[str, Any]: await self._ensure_initialized() return await self._client.wait_processed(timeout=timeout) - async def build_index( - self, - resource_uris: Union[str, List[str]], - **kwargs - ) -> Dict[str, Any]: + async def build_index(self, resource_uris: Union[str, List[str]], **kwargs) -> Dict[str, Any]: """ Manually trigger index building for resources. - + Args: resource_uris: Single URI or list of URIs to index. """ await self._ensure_initialized() return await self._client.build_index(resource_uris, **kwargs) - async def summarize( - self, - resource_uris: Union[str, List[str]], - **kwargs - ) -> Dict[str, Any]: + async def summarize(self, resource_uris: Union[str, List[str]], **kwargs) -> Dict[str, Any]: """ Manually trigger summarization for resources. - + Args: resource_uris: Single URI or list of URIs to summarize. """ diff --git a/openviking/client/local.py b/openviking/client/local.py index 61cefcf9..b437f446 100644 --- a/openviking/client/local.py +++ b/openviking/client/local.py @@ -93,21 +93,13 @@ async def wait_processed(self, timeout: Optional[float] = None) -> Dict[str, Any """Wait for all processing to complete.""" return await self._service.resources.wait_processed(timeout=timeout) - async def build_index( - self, - resource_uris: Union[str, List[str]], - **kwargs - ) -> Dict[str, Any]: + async def build_index(self, resource_uris: Union[str, List[str]], **kwargs) -> Dict[str, Any]: """Manually trigger index building.""" if isinstance(resource_uris, str): resource_uris = [resource_uris] return await self._service.resources.build_index(resource_uris, ctx=self._ctx, **kwargs) - async def summarize( - self, - resource_uris: Union[str, List[str]], - **kwargs - ) -> Dict[str, Any]: + async def summarize(self, resource_uris: Union[str, List[str]], **kwargs) -> Dict[str, Any]: """Manually trigger summarization.""" if isinstance(resource_uris, str): resource_uris = [resource_uris] diff --git a/openviking/eval/ragas/__init__.py b/openviking/eval/ragas/__init__.py index 50032007..03336bc7 100644 --- a/openviking/eval/ragas/__init__.py +++ b/openviking/eval/ragas/__init__.py @@ -47,6 +47,7 @@ class RagasConfig: show_progress: Whether to show progress bar during evaluation. raise_exceptions: Whether to raise exceptions during evaluation. """ + max_workers: int = 16 batch_size: int = 10 timeout: int = 180 @@ -316,12 +317,7 @@ async def evaluate_dataset(self, dataset: EvalDataset) -> SummaryResult: if metric_name in df.columns: scores[metric_name] = float(df.iloc[i][metric_name]) - eval_results.append( - EvalResult( - sample=sample, - scores=scores - ) - ) + eval_results.append(EvalResult(sample=sample, scores=scores)) mean_scores = {} for metric in self.metrics: diff --git a/openviking/eval/ragas/base.py b/openviking/eval/ragas/base.py index c07b4ea4..17f0fca2 100644 --- a/openviking/eval/ragas/base.py +++ b/openviking/eval/ragas/base.py @@ -46,12 +46,7 @@ async def evaluate_dataset(self, dataset: EvalDataset) -> SummaryResult: def _summarize(self, name: str, results: List[EvalResult]) -> SummaryResult: """Aggregate results into a summary.""" if not results: - return SummaryResult( - dataset_name=name, - sample_count=0, - mean_scores={}, - results=[] - ) + return SummaryResult(dataset_name=name, sample_count=0, mean_scores={}, results=[]) metric_sums: Dict[str, float] = {} for res in results: @@ -62,8 +57,5 @@ def _summarize(self, name: str, results: List[EvalResult]) -> SummaryResult: mean_scores = {m: s / count for m, s in metric_sums.items()} return SummaryResult( - dataset_name=name, - sample_count=count, - mean_scores=mean_scores, - results=results + dataset_name=name, sample_count=count, mean_scores=mean_scores, results=results ) diff --git a/openviking/eval/ragas/generator.py b/openviking/eval/ragas/generator.py index 8e46b697..7acfcc91 100644 --- a/openviking/eval/ragas/generator.py +++ b/openviking/eval/ragas/generator.py @@ -69,9 +69,7 @@ async def generate_from_viking_path( # 2. Split content into chunks if needed # 3. Use LLM to generate (Question, Answer, Context) triples return EvalDataset( - name=f"gen_{uuid.uuid4().hex[:8]}", - description=f"Generated from {uri_base}", - samples=[] + name=f"gen_{uuid.uuid4().hex[:8]}", description=f"Generated from {uri_base}", samples=[] ) async def generate_from_content( @@ -127,13 +125,10 @@ async def generate_from_content( query=item["question"], ground_truth=item["answer"], context=[item["context"]], - meta={"source": source_name} + meta={"source": source_name}, ) ) except Exception as e: logger.error(f"Failed to generate samples: {e}") - return EvalDataset( - name=f"gen_{source_name}", - samples=samples - ) + return EvalDataset(name=f"gen_{source_name}", samples=samples) diff --git a/openviking/eval/ragas/pipeline.py b/openviking/eval/ragas/pipeline.py index 11159e48..0b877778 100644 --- a/openviking/eval/ragas/pipeline.py +++ b/openviking/eval/ragas/pipeline.py @@ -157,7 +157,9 @@ def query( if search_result and "results" in search_result: for item in search_result["results"]: uri = item.get("uri", "") - content = item.get("content", "") or item.get("overview", "") or item.get("abstract", "") + content = ( + item.get("content", "") or item.get("overview", "") or item.get("abstract", "") + ) if content: contexts.append(content) retrieved_uris.append(uri) diff --git a/openviking/eval/ragas/play_recorder.py b/openviking/eval/ragas/play_recorder.py index 32919e8c..25f81125 100644 --- a/openviking/eval/ragas/play_recorder.py +++ b/openviking/eval/ragas/play_recorder.py @@ -33,9 +33,6 @@ logger = get_logger(__name__) - - - def print_playback_stats(stats: PlaybackStats) -> None: """Print playback statistics.""" print(f"\n{'=' * 60}") @@ -45,7 +42,11 @@ def print_playback_stats(stats: PlaybackStats) -> None: print(f"\nTotal Records: {stats.total_records}") print(f"Successful: {stats.success_count}") print(f"Failed: {stats.error_count}") - print(f"Success Rate: {stats.success_count / stats.total_records * 100:.1f}%" if stats.total_records > 0 else "N/A") + print( + f"Success Rate: {stats.success_count / stats.total_records * 100:.1f}%" + if stats.total_records > 0 + else "N/A" + ) print("\nPerformance:") print(f" Original Total Latency: {stats.total_original_latency_ms:.2f} ms") @@ -56,7 +57,7 @@ def print_playback_stats(stats: PlaybackStats) -> None: if speedup > 1: print(f" Speedup: {speedup:.2f}x (playback is faster)") else: - print(f" Slowdown: {1/speedup:.2f}x (playback is slower)") + print(f" Slowdown: {1 / speedup:.2f}x (playback is slower)") if stats.total_viking_fs_operations > 0: stats_dict = stats.to_dict() @@ -66,7 +67,9 @@ def print_playback_stats(stats: PlaybackStats) -> None: print("\nVikingFS Detailed Stats:") print(f" Total VikingFS Operations: {viking_fs_stats.get('total_operations', 0)}") print(f" VikingFS Success Rate: {viking_fs_stats.get('success_rate_percent', 0):.1f}%") - print(f" Average AGFS Calls per VikingFS Operation: {viking_fs_stats.get('avg_agfs_calls_per_operation', 0):.2f}") + print( + f" Average AGFS Calls per VikingFS Operation: {viking_fs_stats.get('avg_agfs_calls_per_operation', 0):.2f}" + ) print("\nAGFS FS Detailed Stats:") print(f" Total AGFS Calls: {agfs_fs_stats.get('total_calls', 0)}") diff --git a/openviking/eval/ragas/rag_eval.py b/openviking/eval/ragas/rag_eval.py index b3441def..cab49894 100644 --- a/openviking/eval/ragas/rag_eval.py +++ b/openviking/eval/ragas/rag_eval.py @@ -86,6 +86,7 @@ def __init__( if enable_recorder: from openviking.eval.recorder import init_recorder + init_recorder(enabled=True) logger.info("IO Recorder enabled") @@ -174,11 +175,13 @@ async def retrieve(self, query: str, top_k: int = 5) -> Dict[str, Any]: if result: for ctx in result: - contexts.append({ - "uri": getattr(ctx, "uri", ""), - "content": getattr(ctx, "abstract", "") or getattr(ctx, "overview", ""), - "score": getattr(ctx, "score", 0.0), - }) + contexts.append( + { + "uri": getattr(ctx, "uri", ""), + "content": getattr(ctx, "abstract", "") or getattr(ctx, "overview", ""), + "score": getattr(ctx, "score", 0.0), + } + ) retrieval_time = time.time() - start_time return { @@ -284,9 +287,9 @@ def print_report(eval_results: Dict[str, Any]): for i, result in enumerate(eval_results.get("results", []), 1): print(f"\n[Q{i}] {result['question'][:80]}...") print(f" Contexts Retrieved: {result['context_count']}") - print(f" Retrieval Time: {result['retrieval_time']*1000:.1f}ms") - if result['contexts']: - for j, ctx in enumerate(result['contexts'][:2], 1): + print(f" Retrieval Time: {result['retrieval_time'] * 1000:.1f}ms") + if result["contexts"]: + for j, ctx in enumerate(result["contexts"][:2], 1): print(f" [{j}] URI: {ctx['uri'][:60]}...") print(f" Score: {ctx['score']:.3f}") @@ -383,7 +386,7 @@ async def main_async(args): recorder = get_recorder() viking_fs = get_viking_fs() - if hasattr(viking_fs.agfs, 'stop_recording'): + if hasattr(viking_fs.agfs, "stop_recording"): viking_fs.agfs.stop_recording() stats = recorder.get_stats() @@ -395,10 +398,10 @@ async def main_async(args): print(f"VikingDB Operations: {stats['vikingdb_count']}") print(f"Total Latency: {stats['total_latency_ms']:.2f} ms") print(f"Errors: {stats['errors']}") - if stats['operations']: + if stats["operations"]: print("\nOperations Breakdown:") - for op, data in stats['operations'].items(): - avg_latency = data['total_latency_ms'] / data['count'] if data['count'] > 0 else 0 + for op, data in stats["operations"].items(): + avg_latency = data["total_latency_ms"] / data["count"] if data["count"] > 0 else 0 print(f" {op}: {data['count']} calls, avg {avg_latency:.2f} ms") print(f"\nRecord file: {recorder.record_file}") diff --git a/openviking/eval/ragas/record_analysis.py b/openviking/eval/ragas/record_analysis.py index e4070327..63b5ba97 100644 --- a/openviking/eval/ragas/record_analysis.py +++ b/openviking/eval/ragas/record_analysis.py @@ -222,9 +222,7 @@ def _finalize_operation_stats(stats_dict: Dict[str, OperationStats]) -> None: for stats in stats_dict.values(): if stats.count > 0: stats.avg_latency_ms = stats.total_latency_ms / stats.count - stats.success_rate_percent = ( - stats.success_count / stats.count * 100 - ) + stats.success_rate_percent = stats.success_count / stats.count * 100 else: stats.avg_latency_ms = 0.0 stats.min_latency_ms = 0.0 @@ -308,9 +306,7 @@ def analyze_records( agfs_total = viking_fs_stats.agfs_success_count + viking_fs_stats.agfs_error_count if agfs_total > 0: - viking_fs_stats.agfs_avg_latency_ms = ( - viking_fs_stats.agfs_total_latency_ms / agfs_total - ) + viking_fs_stats.agfs_avg_latency_ms = viking_fs_stats.agfs_total_latency_ms / agfs_total viking_fs_stats.agfs_success_rate_percent = ( viking_fs_stats.agfs_success_count / agfs_total * 100 ) diff --git a/openviking/eval/recorder/playback.py b/openviking/eval/recorder/playback.py index 1573bb6b..9a6155c8 100644 --- a/openviking/eval/recorder/playback.py +++ b/openviking/eval/recorder/playback.py @@ -22,5 +22,6 @@ def get_record_stats(record_file: str) -> dict: stacklevel=2, ) from openviking.eval.record_analysis import analyze_records + stats = analyze_records(record_file) return stats.to_dict() diff --git a/openviking/models/embedder/jina_embedders.py b/openviking/models/embedder/jina_embedders.py index d0c9b396..12f45e4d 100644 --- a/openviking/models/embedder/jina_embedders.py +++ b/openviking/models/embedder/jina_embedders.py @@ -166,4 +166,3 @@ def get_dimension(self) -> int: int: Vector dimension """ return self._dimension - diff --git a/openviking/parse/parsers/code/ast/extractor.py b/openviking/parse/parsers/code/ast/extractor.py index 514756da..46c60f45 100644 --- a/openviking/parse/parsers/code/ast/extractor.py +++ b/openviking/parse/parsers/code/ast/extractor.py @@ -1,4 +1,4 @@ - # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ASTExtractor: language detection + dispatch to per-language extractors.""" @@ -33,8 +33,16 @@ # Language key → (module path, class name, constructor kwargs) _EXTRACTOR_REGISTRY: Dict[str, tuple] = { "python": ("openviking.parse.parsers.code.ast.languages.python", "PythonExtractor", {}), - "javascript": ("openviking.parse.parsers.code.ast.languages.js_ts", "JsTsExtractor", {"lang": "javascript"}), - "typescript": ("openviking.parse.parsers.code.ast.languages.js_ts", "JsTsExtractor", {"lang": "typescript"}), + "javascript": ( + "openviking.parse.parsers.code.ast.languages.js_ts", + "JsTsExtractor", + {"lang": "javascript"}, + ), + "typescript": ( + "openviking.parse.parsers.code.ast.languages.js_ts", + "JsTsExtractor", + {"lang": "typescript"}, + ), "java": ("openviking.parse.parsers.code.ast.languages.java", "JavaExtractor", {}), "cpp": ("openviking.parse.parsers.code.ast.languages.cpp", "CppExtractor", {}), "rust": ("openviking.parse.parsers.code.ast.languages.rust", "RustExtractor", {}), @@ -71,11 +79,15 @@ def _get_extractor(self, lang: Optional[str]) -> Optional[LanguageExtractor]: self._cache[lang] = extractor return extractor except Exception as e: - logger.warning("AST extractor unavailable for language '%s', falling back to LLM: %s", lang, e) + logger.warning( + "AST extractor unavailable for language '%s', falling back to LLM: %s", lang, e + ) self._cache[lang] = None return None - def extract_skeleton(self, file_name: str, content: str, verbose: bool = False) -> Optional[str]: + def extract_skeleton( + self, file_name: str, content: str, verbose: bool = False + ) -> Optional[str]: """Extract skeleton text from source code. Returns None for unsupported languages or on extraction failure, @@ -94,7 +106,12 @@ def extract_skeleton(self, file_name: str, content: str, verbose: bool = False) skeleton: CodeSkeleton = extractor.extract(file_name, content) return skeleton.to_text(verbose=verbose) except Exception as e: - logger.warning("AST extraction failed for '%s' (language: %s), falling back to LLM: %s", file_name, lang, e) + logger.warning( + "AST extraction failed for '%s' (language: %s), falling back to LLM: %s", + file_name, + lang, + e, + ) return None diff --git a/openviking/parse/parsers/code/ast/languages/cpp.py b/openviking/parse/parsers/code/ast/languages/cpp.py index 305d529a..b3237675 100644 --- a/openviking/parse/parsers/code/ast/languages/cpp.py +++ b/openviking/parse/parsers/code/ast/languages/cpp.py @@ -9,7 +9,7 @@ def _node_text(node, content_bytes: bytes) -> str: - return content_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace") + return content_bytes[node.start_byte : node.end_byte].decode("utf-8", errors="replace") def _parse_block_comment(raw: str) -> str: @@ -65,8 +65,13 @@ def _extract_function(node, content_bytes: bytes, docstring: str = "") -> Functi for child in node.children: if child.type == "function_declarator": name, params = _extract_function_declarator(child, content_bytes) - elif child.type in ("type_specifier", "primitive_type", "type_identifier", - "qualified_identifier", "auto"): + elif child.type in ( + "type_specifier", + "primitive_type", + "type_identifier", + "qualified_identifier", + "auto", + ): if not return_type: return_type = _node_text(child, content_bytes) elif child.type == "pointer_declarator": @@ -104,15 +109,27 @@ def _extract_class(node, content_bytes: bytes, docstring: str = "") -> ClassSkel fn_name = "" fn_params = "" for sub in child.children: - if sub.type in ("type_specifier", "primitive_type", "type_identifier", - "qualified_identifier") and not ret_type: + if ( + sub.type + in ( + "type_specifier", + "primitive_type", + "type_identifier", + "qualified_identifier", + ) + and not ret_type + ): ret_type = _node_text(sub, content_bytes) elif sub.type == "function_declarator": fn_name, fn_params = _extract_function_declarator(sub, content_bytes) break if fn_name: doc = _preceding_doc(siblings, idx, content_bytes) - methods.append(FunctionSig(name=fn_name, params=fn_params, return_type=ret_type, docstring=doc)) + methods.append( + FunctionSig( + name=fn_name, params=fn_params, return_type=ret_type, docstring=doc + ) + ) return ClassSkeleton(name=name, bases=bases, docstring=docstring, methods=methods) @@ -157,7 +174,9 @@ def extract(self, file_name: str, content: str) -> CodeSkeleton: classes.append(_extract_class(s2, content_bytes, docstring=doc)) elif s2.type == "function_definition": doc = _preceding_doc(inner, i2, content_bytes) - functions.append(_extract_function(s2, content_bytes, docstring=doc)) + functions.append( + _extract_function(s2, content_bytes, docstring=doc) + ) return CodeSkeleton( file_name=file_name, diff --git a/openviking/parse/parsers/code/ast/languages/go.py b/openviking/parse/parsers/code/ast/languages/go.py index a067a7fb..d3ca44f0 100644 --- a/openviking/parse/parsers/code/ast/languages/go.py +++ b/openviking/parse/parsers/code/ast/languages/go.py @@ -9,7 +9,7 @@ def _node_text(node, content_bytes: bytes) -> str: - return content_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace") + return content_bytes[node.start_byte : node.end_byte].decode("utf-8", errors="replace") def _preceding_doc(siblings: list, idx: int, content_bytes: bytes) -> str: @@ -92,7 +92,9 @@ def extract(self, file_name: str, content: str) -> CodeSkeleton: if s2.type == "import_spec": for s3 in s2.children: if s3.type == "interpreted_string_literal": - imports.append(_node_text(s3, content_bytes).strip().strip('"')) + imports.append( + _node_text(s3, content_bytes).strip().strip('"') + ) elif child.type in ("function_declaration", "method_declaration"): doc = _preceding_doc(siblings, idx, content_bytes) functions.append(_extract_function(child, content_bytes, docstring=doc)) diff --git a/openviking/parse/parsers/code/ast/languages/java.py b/openviking/parse/parsers/code/ast/languages/java.py index 55103c40..79de06dc 100644 --- a/openviking/parse/parsers/code/ast/languages/java.py +++ b/openviking/parse/parsers/code/ast/languages/java.py @@ -9,7 +9,7 @@ def _node_text(node, content_bytes: bytes) -> str: - return content_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace") + return content_bytes[node.start_byte : node.end_byte].decode("utf-8", errors="replace") def _parse_block_comment(raw: str) -> str: @@ -44,9 +44,15 @@ def _extract_method(node, content_bytes: bytes, docstring: str = "") -> Function if child.type == "identifier" and not name: if return_type: name = _node_text(child, content_bytes) - elif child.type in ("type_identifier", "void_type", "integral_type", - "floating_point_type", "boolean_type", "array_type", - "generic_type"): + elif child.type in ( + "type_identifier", + "void_type", + "integral_type", + "floating_point_type", + "boolean_type", + "array_type", + "generic_type", + ): if not return_type: return_type = _node_text(child, content_bytes) elif child.type == "formal_parameters": diff --git a/openviking/parse/parsers/code/ast/languages/js_ts.py b/openviking/parse/parsers/code/ast/languages/js_ts.py index a092f7dc..d0422682 100644 --- a/openviking/parse/parsers/code/ast/languages/js_ts.py +++ b/openviking/parse/parsers/code/ast/languages/js_ts.py @@ -9,7 +9,7 @@ def _node_text(node, content_bytes: bytes) -> str: - return content_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace") + return content_bytes[node.start_byte : node.end_byte].decode("utf-8", errors="replace") def _parse_jsdoc(raw: str) -> str: @@ -49,7 +49,7 @@ def _first_string_in_body(body_node, content_bytes: bytes) -> str: raw = _node_text(sub, content_bytes).strip() for q in ('"""', "'''", '"', "'", "`"): if raw.startswith(q) and raw.endswith(q) and len(raw) >= 2 * len(q): - return raw[len(q):-len(q)].strip() + return raw[len(q) : -len(q)].strip() return raw break return "" @@ -64,7 +64,9 @@ def _extract_params(params_node, content_bytes: bytes) -> str: return raw.strip() -def _extract_function(node, content_bytes: bytes, lang_name: str, docstring: str = "") -> FunctionSig: +def _extract_function( + node, content_bytes: bytes, lang_name: str, docstring: str = "" +) -> FunctionSig: name = "" params = "" return_type = "" @@ -80,7 +82,7 @@ def _extract_function(node, content_bytes: bytes, lang_name: str, docstring: str elif child.type == "type_annotation": # TypeScript return type annotation for sub in child.children: - if sub.type not in (":", ): + if sub.type not in (":",): return_type = _node_text(sub, content_bytes).strip() break elif child.type == "statement_block": @@ -91,7 +93,9 @@ def _extract_function(node, content_bytes: bytes, lang_name: str, docstring: str return FunctionSig(name=name, params=params, return_type=return_type, docstring=docstring) -def _extract_class(node, content_bytes: bytes, lang_name: str, docstring: str = "") -> ClassSkeleton: +def _extract_class( + node, content_bytes: bytes, lang_name: str, docstring: str = "" +) -> ClassSkeleton: name = "" bases: List[str] = [] body_node = None @@ -144,12 +148,15 @@ def __init__(self, lang: str): if lang == "javascript": import tree_sitter_javascript as tsjs from tree_sitter import Language, Parser + self._language = Language(tsjs.language()) else: import tree_sitter_typescript as tsts from tree_sitter import Language, Parser + self._language = Language(tsts.language_typescript()) from tree_sitter import Parser + self._parser = Parser(self._language) def extract(self, file_name: str, content: str) -> CodeSkeleton: @@ -168,7 +175,7 @@ def extract(self, file_name: str, content: str) -> CodeSkeleton: # from "module" for sub in child.children: if sub.type == "string": - raw = _node_text(sub, content_bytes).strip().strip('"\'') + raw = _node_text(sub, content_bytes).strip().strip("\"'") if raw not in _seen_imports: imports.append(raw) _seen_imports.add(raw) @@ -178,17 +185,23 @@ def extract(self, file_name: str, content: str) -> CodeSkeleton: classes.append(_extract_class(child, content_bytes, self._lang_name, docstring=doc)) elif child.type in ("function_declaration", "generator_function_declaration"): doc = _preceding_doc(siblings, idx, content_bytes) - functions.append(_extract_function(child, content_bytes, self._lang_name, docstring=doc)) + functions.append( + _extract_function(child, content_bytes, self._lang_name, docstring=doc) + ) elif child.type == "export_statement": # export default class / export function for sub in child.children: if sub.type == "class_declaration": doc = _preceding_doc(siblings, idx, content_bytes) - classes.append(_extract_class(sub, content_bytes, self._lang_name, docstring=doc)) + classes.append( + _extract_class(sub, content_bytes, self._lang_name, docstring=doc) + ) break elif sub.type in ("function_declaration", "generator_function_declaration"): doc = _preceding_doc(siblings, idx, content_bytes) - functions.append(_extract_function(sub, content_bytes, self._lang_name, docstring=doc)) + functions.append( + _extract_function(sub, content_bytes, self._lang_name, docstring=doc) + ) break elif child.type == "lexical_declaration": # const foo = () => ... @@ -200,7 +213,9 @@ def extract(self, file_name: str, content: str) -> CodeSkeleton: fn_name = _node_text(s2, content_bytes) elif s2.type == "arrow_function": doc = _preceding_doc(siblings, idx, content_bytes) - fn = _extract_function(s2, content_bytes, self._lang_name, docstring=doc) + fn = _extract_function( + s2, content_bytes, self._lang_name, docstring=doc + ) fn.name = fn_name functions.append(fn) diff --git a/openviking/parse/parsers/code/ast/languages/python.py b/openviking/parse/parsers/code/ast/languages/python.py index caa74c15..b1636972 100644 --- a/openviking/parse/parsers/code/ast/languages/python.py +++ b/openviking/parse/parsers/code/ast/languages/python.py @@ -9,7 +9,7 @@ def _node_text(node, content_bytes: bytes) -> str: - return content_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace") + return content_bytes[node.start_byte : node.end_byte].decode("utf-8", errors="replace") def _first_string_child(body_node, content_bytes: bytes) -> str: @@ -24,7 +24,7 @@ def _first_string_child(body_node, content_bytes: bytes) -> str: # Strip quotes for q in ('"""', "'''", '"', "'"): if raw.startswith(q) and raw.endswith(q) and len(raw) >= 2 * len(q): - return raw[len(q):-len(q)].strip() + return raw[len(q) : -len(q)].strip() return raw break # only check first expression_statement return "" @@ -153,7 +153,7 @@ def extract(self, file_name: str, content: str) -> CodeSkeleton: raw = _node_text(sub, content_bytes).strip() for q in ('"""', "'''", '"', "'"): if raw.startswith(q) and raw.endswith(q) and len(raw) >= 2 * len(q): - module_doc = raw[len(q):-len(q)].strip() + module_doc = raw[len(q) : -len(q)].strip() break else: module_doc = raw diff --git a/openviking/parse/parsers/code/ast/languages/rust.py b/openviking/parse/parsers/code/ast/languages/rust.py index 9479c2b6..953d25b8 100644 --- a/openviking/parse/parsers/code/ast/languages/rust.py +++ b/openviking/parse/parsers/code/ast/languages/rust.py @@ -9,7 +9,7 @@ def _node_text(node, content_bytes: bytes) -> str: - return content_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace") + return content_bytes[node.start_byte : node.end_byte].decode("utf-8", errors="replace") def _preceding_doc(siblings: list, idx: int, content_bytes: bytes) -> str: @@ -103,7 +103,9 @@ def extract(self, file_name: str, content: str) -> CodeSkeleton: siblings = list(root.children) for idx, child in enumerate(siblings): if child.type == "use_declaration": - imports.append(_node_text(child, content_bytes).strip().rstrip(";").replace("use ", "")) + imports.append( + _node_text(child, content_bytes).strip().rstrip(";").replace("use ", "") + ) elif child.type in ("struct_item", "trait_item", "enum_item"): doc = _preceding_doc(siblings, idx, content_bytes) classes.append(_extract_struct_or_trait(child, content_bytes, docstring=doc)) diff --git a/openviking/parse/parsers/code/ast/skeleton.py b/openviking/parse/parsers/code/ast/skeleton.py index cb4bb069..110e882a 100644 --- a/openviking/parse/parsers/code/ast/skeleton.py +++ b/openviking/parse/parsers/code/ast/skeleton.py @@ -15,9 +15,9 @@ def _compact_params(params: str) -> str: @dataclass class FunctionSig: name: str - params: str # raw parameter string, e.g. "source, instruction, **kwargs" + params: str # raw parameter string, e.g. "source, instruction, **kwargs" return_type: str # e.g. "ParseResult" or "" - docstring: str # first line only + docstring: str # first line only @dataclass @@ -33,9 +33,9 @@ class CodeSkeleton: file_name: str language: str module_doc: str - imports: List[str] # flattened, e.g. ["asyncio", "os", "typing.Optional"] + imports: List[str] # flattened, e.g. ["asyncio", "os", "typing.Optional"] classes: List[ClassSkeleton] - functions: List[FunctionSig] # top-level functions only + functions: List[FunctionSig] # top-level functions only def to_text(self, verbose: bool = False) -> str: """Generate skeleton text. @@ -44,6 +44,7 @@ def to_text(self, verbose: bool = False) -> str: verbose: If True, include full docstrings (for ast_llm mode / LLM input). If False, only keep the first line (for ast mode / direct embedding). """ + def _doc(raw: str, indent: str) -> List[str]: if not raw: return [] @@ -54,9 +55,11 @@ def _doc(raw: str, indent: str) -> List[str]: doc_lines = raw.strip().split("\n") if len(doc_lines) == 1: return [f'{indent}"""{first}"""'] - return [f'{indent}"""{doc_lines[0]}'] + \ - [f'{indent}{l.strip()}' for l in doc_lines[1:]] + \ - [f'{indent}"""'] + return ( + [f'{indent}"""{doc_lines[0]}'] + + [f"{indent}{l.strip()}" for l in doc_lines[1:]] + + [f'{indent}"""'] + ) lines: List[str] = [] diff --git a/openviking/parse/parsers/code/code.py b/openviking/parse/parsers/code/code.py index ea6504bb..3bd06d7c 100644 --- a/openviking/parse/parsers/code/code.py +++ b/openviking/parse/parsers/code/code.py @@ -311,11 +311,15 @@ async def _run_git(self, args: List[str], cwd: Optional[str] = None) -> str: error_msg = stderr.decode().strip() user_msg = "Git command failed." if "Could not resolve hostname" in error_msg: - user_msg = "Git command failed: could not resolve hostname. Check the URL or your network." + user_msg = ( + "Git command failed: could not resolve hostname. Check the URL or your network." + ) elif "Permission denied" in error_msg or "publickey" in error_msg: - user_msg = "Git command failed: authentication error. Check your SSH keys or credentials." + user_msg = ( + "Git command failed: authentication error. Check your SSH keys or credentials." + ) raise RuntimeError( - f'{user_msg} Command: git {" ".join(args[1:])}. Details: {error_msg}' + f"{user_msg} Command: git {' '.join(args[1:])}. Details: {error_msg}" ) return stdout.decode().strip() diff --git a/openviking/parse/tree_builder.py b/openviking/parse/tree_builder.py index 820ca554..a6ce29dd 100644 --- a/openviking/parse/tree_builder.py +++ b/openviking/parse/tree_builder.py @@ -180,7 +180,9 @@ async def finalize_from_temp( await self._enqueue_semantic_generation(final_uri, "resource", ctx=ctx) logger.info(f"[TreeBuilder] Enqueued semantic generation for: {final_uri}") except Exception as e: - logger.error(f"[TreeBuilder] Failed to enqueue semantic generation: {e}", exc_info=True) + logger.error( + f"[TreeBuilder] Failed to enqueue semantic generation: {e}", exc_info=True + ) # 7. Return simple BuildingTree (no scanning needed) tree = BuildingTree( @@ -188,7 +190,7 @@ async def finalize_from_temp( source_format=source_format, ) tree._root_uri = final_uri - + # Create a minimal Context object for the root so that tree.root is not None root_context = Context(uri=final_uri) tree.add_context(root_context) diff --git a/openviking/server/bootstrap.py b/openviking/server/bootstrap.py index 942c3c91..ed2ea3b6 100644 --- a/openviking/server/bootstrap.py +++ b/openviking/server/bootstrap.py @@ -159,9 +159,7 @@ def _start_vikingbot_gateway(enable_logging: bool, log_dir: str) -> Optional[Bot python_cmd = sys.executable try: result = subprocess.run( - [python_cmd, "-m", "vikingbot", "--help"], - capture_output=True, - timeout=5 + [python_cmd, "-m", "vikingbot", "--help"], capture_output=True, timeout=5 ) if result.returncode == 0: vikingbot_cmd = [python_cmd, "-m", "vikingbot", "gateway"] diff --git a/openviking/server/routers/bot.py b/openviking/server/routers/bot.py index bc819a18..f93742cd 100644 --- a/openviking/server/routers/bot.py +++ b/openviking/server/routers/bot.py @@ -64,7 +64,7 @@ async def health_check(request: Request): try: async with httpx.AsyncClient() as client: - print(f'url={f"{bot_url}/bot/v1/health"}') + print(f"url={f'{bot_url}/bot/v1/health'}") # Forward to Vikingbot OpenAPIChannel health endpoint response = await client.get( f"{bot_url}/bot/v1/health", diff --git a/openviking/service/resource_service.py b/openviking/service/resource_service.py index 54a4e2cc..39612ab2 100644 --- a/openviking/service/resource_service.py +++ b/openviking/service/resource_service.py @@ -173,10 +173,7 @@ async def add_skill( return result async def build_index( - self, - resource_uris: List[str], - ctx: RequestContext, - **kwargs + self, resource_uris: List[str], ctx: RequestContext, **kwargs ) -> Dict[str, Any]: """Manually trigger index building. @@ -191,10 +188,7 @@ async def build_index( return await self._resource_processor.build_index(resource_uris, ctx, **kwargs) async def summarize( - self, - resource_uris: List[str], - ctx: RequestContext, - **kwargs + self, resource_uris: List[str], ctx: RequestContext, **kwargs ) -> Dict[str, Any]: """Manually trigger summarization. diff --git a/openviking/session/__init__.py b/openviking/session/__init__.py index d238cc35..6add1ab9 100644 --- a/openviking/session/__init__.py +++ b/openviking/session/__init__.py @@ -10,7 +10,12 @@ MemoryActionDecision, MemoryDeduplicator, ) -from openviking.session.memory_extractor import CandidateMemory, MemoryCategory, MemoryExtractor, ToolSkillCandidateMemory +from openviking.session.memory_extractor import ( + CandidateMemory, + MemoryCategory, + MemoryExtractor, + ToolSkillCandidateMemory, +) from openviking.session.session import Session, SessionCompression, SessionStats __all__ = [ diff --git a/openviking/storage/vectordb/project/vikingdb_project.py b/openviking/storage/vectordb/project/vikingdb_project.py index bcf82ec9..3f3d3bfd 100644 --- a/openviking/storage/vectordb/project/vikingdb_project.py +++ b/openviking/storage/vectordb/project/vikingdb_project.py @@ -124,7 +124,7 @@ def get_collection(self, collection_name: str) -> Optional[Collection]: } # Update with user-provided arguments (can override defaults if needed, though usually additive) kwargs.update(self.collection_args) - + vikingdb_collection = self.CollectionClass(**kwargs) return Collection(vikingdb_collection) except Exception: @@ -154,23 +154,21 @@ def list_collections(self) -> List[str]: def get_collections(self) -> Dict[str, Collection]: """Get all collections from server""" colls = self._get_collections() - + # Prepare base arguments base_kwargs = { "host": self.host, "headers": self.headers, } - + collections = {} for c in colls: kwargs = base_kwargs.copy() kwargs["meta_data"] = c kwargs.update(self.collection_args) - - collections[c["CollectionName"]] = Collection( - self.CollectionClass(**kwargs) - ) - + + collections[c["CollectionName"]] = Collection(self.CollectionClass(**kwargs)) + return collections def create_collection(self, collection_name: str, meta_data: Dict[str, Any]) -> Collection: diff --git a/openviking/storage/vectordb_adapters/factory.py b/openviking/storage/vectordb_adapters/factory.py index 62039293..6c11954c 100644 --- a/openviking/storage/vectordb_adapters/factory.py +++ b/openviking/storage/vectordb_adapters/factory.py @@ -27,6 +27,7 @@ def create_collection_adapter(config) -> CollectionAdapter: if adapter_cls is None and "." in backend: try: import importlib + module_name, class_name = backend.rsplit(".", 1) module = importlib.import_module(module_name) potential_cls = getattr(module, class_name) diff --git a/openviking/storage/viking_fs.py b/openviking/storage/viking_fs.py index 24e8da5f..46e2f57f 100644 --- a/openviking/storage/viking_fs.py +++ b/openviking/storage/viking_fs.py @@ -306,9 +306,8 @@ async def grep( """Content search by pattern or keywords.""" self._ensure_access(uri, ctx) path = self._uri_to_path(uri, ctx=ctx) - result = await asyncio.to_thread( - self.agfs.grep, - path, pattern, True, case_insensitive, False, node_limit=node_limit + result = await asyncio.to_thread( + self.agfs.grep, path, pattern, True, case_insensitive, False, node_limit=node_limit ) if result.get("matches", None) is None: result["matches"] = [] diff --git a/openviking/utils/embedding_utils.py b/openviking/utils/embedding_utils.py index 5f653065..ca2ba36a 100644 --- a/openviking/utils/embedding_utils.py +++ b/openviking/utils/embedding_utils.py @@ -20,6 +20,7 @@ logger = get_logger(__name__) + def _owner_space_for_uri(uri: str, ctx: RequestContext) -> str: """Derive owner_space from a URI.""" if uri.startswith("viking://agent/"): @@ -28,15 +29,31 @@ def _owner_space_for_uri(uri: str, ctx: RequestContext) -> str: return ctx.user.user_space_name() return "" + def get_resource_content_type(file_name: str) -> ResourceContentType: """Determine resource content type based on file extension.""" file_name = file_name.lower() - - text_extensions = {".txt", ".md", ".csv", ".json", ".xml", ".py", ".js", ".ts", ".java", ".cpp", ".c", ".h", ".go", ".rs"} + + text_extensions = { + ".txt", + ".md", + ".csv", + ".json", + ".xml", + ".py", + ".js", + ".ts", + ".java", + ".cpp", + ".c", + ".h", + ".go", + ".rs", + } image_extensions = {".png", ".jpg", ".jpeg", ".gif", ".bmp", ".svg", ".webp"} video_extensions = {".mp4", ".avi", ".mov", ".wmv", ".flv"} audio_extensions = {".mp3", ".wav", ".aac", ".flac"} - + if any(file_name.endswith(ext) for ext in text_extensions): return ResourceContentType.TEXT elif any(file_name.endswith(ext) for ext in image_extensions): @@ -45,9 +62,10 @@ def get_resource_content_type(file_name: str) -> ResourceContentType: return ResourceContentType.VIDEO elif any(file_name.endswith(ext) for ext in audio_extensions): return ResourceContentType.AUDIO - + return ResourceContentType.UNKNOWN + async def vectorize_directory_meta( uri: str, abstract: str, @@ -57,7 +75,7 @@ async def vectorize_directory_meta( ) -> None: """ Vectorize directory metadata (.abstract.md and .overview.md). - + Creates Context objects for abstract and overview and enqueues them. """ if not ctx: @@ -66,7 +84,7 @@ async def vectorize_directory_meta( queue_manager = get_queue_manager() embedding_queue = queue_manager.get_queue(queue_manager.EMBEDDING) - + parent_uri = VikingURI(uri).parent.uri owner_space = _owner_space_for_uri(uri, ctx) @@ -106,6 +124,7 @@ async def vectorize_directory_meta( await embedding_queue.enqueue(msg_overview) logger.debug(f"Enqueued directory L1 (overview) for vectorization: {uri}") + async def vectorize_file( file_path: str, summary_dict: Dict[str, str], @@ -115,7 +134,7 @@ async def vectorize_file( ) -> None: """ Vectorize a single file. - + Creates Context object for the file and enqueues it. Reads content for TEXT files, otherwise uses summary. """ @@ -152,7 +171,9 @@ async def vectorize_file( content = content.decode("utf-8", errors="replace") context.set_vectorize(Vectorize(text=content)) except Exception as e: - logger.warning(f"Failed to read file content for {file_path}, falling back to summary: {e}") + logger.warning( + f"Failed to read file content for {file_path}, falling back to summary: {e}" + ) if summary: context.set_vectorize(Vectorize(text=summary)) else: @@ -168,69 +189,67 @@ async def vectorize_file( embedding_msg = EmbeddingMsgConverter.from_context(context) if not embedding_msg: return - + await embedding_queue.enqueue(embedding_msg) logger.debug(f"Enqueued file for vectorization: {file_path}") - + except Exception as e: logger.error(f"Failed to vectorize file {file_path}: {e}", exc_info=True) + async def index_resource( uri: str, ctx: RequestContext, ) -> None: """ Build vector index for a resource directory. - + 1. Reads .abstract.md and .overview.md and vectorizes them. 2. Scans files in the directory and vectorizes them. """ viking_fs = get_viking_fs() - + # 1. Index Directory Metadata abstract_uri = f"{uri}/.abstract.md" overview_uri = f"{uri}/.overview.md" - + abstract = "" overview = "" - + if await viking_fs.exists(abstract_uri): content = await viking_fs.read_file(abstract_uri) if isinstance(content, bytes): abstract = content.decode("utf-8") - + if await viking_fs.exists(overview_uri): content = await viking_fs.read_file(overview_uri) if isinstance(content, bytes): overview = content.decode("utf-8") - + if abstract or overview: await vectorize_directory_meta(uri, abstract, overview, ctx=ctx) - + # 2. Index Files try: files = await viking_fs.ls(uri, ctx=ctx) for file_info in files: file_name = file_info["name"] - + # Skip hidden files (like .abstract.md) if file_name.startswith("."): continue - + if file_info.get("type") == "directory" or file_info.get("isDir"): # TODO: Recursive indexing? For now, skip subdirectories to match previous behavior continue - + file_uri = file_info.get("uri") or f"{uri}/{file_name}" - + # For direct indexing, we might not have summaries. # We pass empty summary_dict, vectorize_file will try to read content for text files. await vectorize_file( - file_path=file_uri, - summary_dict={"name": file_name}, - parent_uri=uri, - ctx=ctx + file_path=file_uri, summary_dict={"name": file_name}, parent_uri=uri, ctx=ctx ) - + except Exception as e: logger.error(f"Failed to scan directory {uri} for indexing: {e}") diff --git a/openviking/utils/resource_processor.py b/openviking/utils/resource_processor.py index 3ac5beaa..605f8cab 100644 --- a/openviking/utils/resource_processor.py +++ b/openviking/utils/resource_processor.py @@ -78,15 +78,20 @@ def _get_media_processor(self): ) return self._media_processor - async def build_index(self, resource_uris: List[str], ctx: RequestContext, **kwargs) -> Dict[str, Any]: + async def build_index( + self, resource_uris: List[str], ctx: RequestContext, **kwargs + ) -> Dict[str, Any]: """Expose index building as a standalone method.""" for uri in resource_uris: await index_resource(uri, ctx) return {"status": "success", "message": f"Indexed {len(resource_uris)} resources"} - async def summarize(self, resource_uris: List[str], ctx: RequestContext, **kwargs) -> Dict[str, Any]: + async def summarize( + self, resource_uris: List[str], ctx: RequestContext, **kwargs + ) -> Dict[str, Any]: """Expose summarization as a standalone method.""" return await self._get_summarizer().summarize(resource_uris, ctx, **kwargs) + async def process_resource( self, path: str, @@ -195,32 +200,29 @@ async def process_resource( # ============ Phase 4: Optional Steps ============ if summarize: - # Explicit summarization request. - # If build_index is ALSO True, we want vectorization. - # If build_index is False, we skip vectorization. - skip_vec = not build_index - try: + # Explicit summarization request. + # If build_index is ALSO True, we want vectorization. + # If build_index is False, we skip vectorization. + skip_vec = not build_index + try: await self._get_summarizer().summarize( resource_uris=[result["root_uri"]], ctx=ctx, skip_vectorization=skip_vec, - **kwargs + **kwargs, ) - except Exception as e: + except Exception as e: logger.error(f"Summarization failed: {e}") result["warnings"] = result.get("warnings", []) + [f"Summarization failed: {e}"] elif build_index: - # Standard compatibility mode: "Just Index it" usually implies ingestion flow. - # We assume this means "Ingest and Index", which requires summarization. - try: + # Standard compatibility mode: "Just Index it" usually implies ingestion flow. + # We assume this means "Ingest and Index", which requires summarization. + try: await self._get_summarizer().summarize( - resource_uris=[result["root_uri"]], - ctx=ctx, - skip_vectorization=False, - **kwargs + resource_uris=[result["root_uri"]], ctx=ctx, skip_vectorization=False, **kwargs ) - except Exception as e: + except Exception as e: logger.error(f"Auto-index failed: {e}") result["warnings"] = result.get("warnings", []) + [f"Auto-index failed: {e}"] diff --git a/openviking/utils/summarizer.py b/openviking/utils/summarizer.py index 8b3a0939..a7477ba3 100644 --- a/openviking/utils/summarizer.py +++ b/openviking/utils/summarizer.py @@ -16,11 +16,12 @@ logger = get_logger(__name__) + class Summarizer: """ Handles summarization of resources. """ - + def __init__(self, vlm_processor: "VLMProcessor"): self.vlm_processor = vlm_processor @@ -29,7 +30,7 @@ async def summarize( resource_uris: List[str], ctx: "RequestContext", skip_vectorization: bool = False, - **kwargs + **kwargs, ) -> Dict[str, Any]: """ Summarize the given resources. @@ -37,7 +38,7 @@ async def summarize( """ queue_manager = get_queue_manager() semantic_queue = queue_manager.get_queue(queue_manager.SEMANTIC, allow_create=True) - + enqueued_count = 0 for uri in resource_uris: # Determine context_type based on URI @@ -46,7 +47,7 @@ async def summarize( context_type = "memory" elif uri.startswith("viking://agent/skills/"): context_type = "skill" - + msg = SemanticMsg( uri=uri, context_type=context_type, @@ -58,6 +59,8 @@ async def summarize( ) await semantic_queue.enqueue(msg) enqueued_count += 1 - logger.info(f"Enqueued semantic generation for: {uri} (skip_vectorization={skip_vectorization})") - + logger.info( + f"Enqueued semantic generation for: {uri} (skip_vectorization={skip_vectorization})" + ) + return {"status": "success", "enqueued_count": enqueued_count} diff --git a/openviking_cli/utils/config/vectordb_config.py b/openviking_cli/utils/config/vectordb_config.py index 8b4a4e16..6b77747a 100644 --- a/openviking_cli/utils/config/vectordb_config.py +++ b/openviking_cli/utils/config/vectordb_config.py @@ -110,7 +110,7 @@ class VectorDBBackendConfig(BaseModel): def validate_config(self): """Validate configuration completeness and consistency""" standard_backends = ["local", "http", "volcengine", "vikingdb"] - + # Allow custom backend classes (containing dot) without standard validation if "." in self.backend: logger.info("Using custom VectorDB backend: %s", self.backend) @@ -145,4 +145,4 @@ def validate_config(self): if not self.vikingdb or not self.vikingdb.host: raise ValueError("VectorDB vikingdb backend requires 'host' to be set") - return self \ No newline at end of file + return self diff --git a/tests/client/test_import_export.py b/tests/client/test_import_export.py index ec66043a..e4dfe3a9 100644 --- a/tests/client/test_import_export.py +++ b/tests/client/test_import_export.py @@ -116,7 +116,6 @@ async def test_import_export_roundtrip( # Verify content consistency assert original_content == imported_content - @staticmethod def _build_ovpack(zip_path: Path, entries: dict[str, str]) -> None: buffer = io.BytesIO() @@ -172,4 +171,6 @@ async def test_import_rejects_unsafe_entries( self._build_ovpack(ovpack_path, entries) with pytest.raises(ValueError, match=error_pattern): - await client.import_ovpack(str(ovpack_path), "viking://resources/security/", vectorize=False) + await client.import_ovpack( + str(ovpack_path), "viking://resources/security/", vectorize=False + ) diff --git a/tests/integration/test_add_resource_index.py b/tests/integration/test_add_resource_index.py index 69a664c3..32421e69 100644 --- a/tests/integration/test_add_resource_index.py +++ b/tests/integration/test_add_resource_index.py @@ -10,76 +10,67 @@ from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton from tests.utils.mock_agfs import MockLocalAGFS + @pytest.fixture def test_config(tmp_path): """Create a temporary config file.""" config_path = tmp_path / "ov.conf" workspace = tmp_path / "workspace" workspace.mkdir() - + config_content = { "storage": { "workspace": str(workspace), - "agfs": { - "backend": "local", - "port": 1833 - }, - "vectordb": { - "backend": "local" - } + "agfs": {"backend": "local", "port": 1833}, + "vectordb": {"backend": "local"}, }, "embedding": { - "dense": { - "provider": "openai", - "api_key": "fake", - "model": "text-embedding-3-small" - } + "dense": {"provider": "openai", "api_key": "fake", "model": "text-embedding-3-small"} }, - "vlm": { - "provider": "openai", - "api_key": "fake", - "model": "gpt-4-vision-preview" - } + "vlm": {"provider": "openai", "api_key": "fake", "model": "gpt-4-vision-preview"}, } config_path.write_text(json.dumps(config_content)) return config_path + @pytest.fixture async def client(test_config, tmp_path): """Initialize AsyncOpenViking client with mocks.""" - + # Set config env var os.environ["OPENVIKING_CONFIG_FILE"] = str(test_config) - + # Reset Singletons OpenVikingConfigSingleton._instance = None await AsyncOpenViking.reset() - + mock_agfs = MockLocalAGFS(root_path=tmp_path / "mock_agfs_root") # Mock LLM/VLM services AND AGFS - with patch("openviking.utils.summarizer.Summarizer.summarize") as mock_summarize, \ - patch("openviking.utils.index_builder.IndexBuilder.build_index") as mock_build_index, \ - patch("openviking.utils.agfs_utils.create_agfs_client", return_value=mock_agfs), \ - patch("openviking.agfs_manager.AGFSManager.start"), \ - patch("openviking.agfs_manager.AGFSManager.stop"): - + with ( + patch("openviking.utils.summarizer.Summarizer.summarize") as mock_summarize, + patch("openviking.utils.index_builder.IndexBuilder.build_index") as mock_build_index, + patch("openviking.utils.agfs_utils.create_agfs_client", return_value=mock_agfs), + patch("openviking.agfs_manager.AGFSManager.start"), + patch("openviking.agfs_manager.AGFSManager.stop"), + ): # Make mocks return success mock_summarize.return_value = {"status": "success"} mock_build_index.return_value = {"status": "success"} - + client = AsyncOpenViking(path=str(test_config.parent)) await client.initialize() - + yield client - + await client.close() - + # Cleanup OpenVikingConfigSingleton._instance = None if "OPENVIKING_CONFIG_FILE" in os.environ: del os.environ["OPENVIKING_CONFIG_FILE"] + @pytest.mark.asyncio async def test_add_resource_indexing_logic(test_config, tmp_path): """ @@ -94,61 +85,54 @@ async def test_add_resource_indexing_logic(test_config, tmp_path): # Create dummy resource resource_file = tmp_path / "test_doc.md" resource_file.write_text("# Test Document\n\nThis is a test document.", encoding="utf-8") - + mock_agfs = MockLocalAGFS(root_path=tmp_path / "mock_agfs_root") # Patch the Summarizer and IndexBuilder to verify calls - with patch("openviking.utils.summarizer.Summarizer.summarize", new_callable=AsyncMock) as mock_summarize, \ - patch("openviking.utils.agfs_utils.create_agfs_client", return_value=mock_agfs), \ - patch("openviking.agfs_manager.AGFSManager.start"), \ - patch("openviking.agfs_manager.AGFSManager.stop"): - + with ( + patch( + "openviking.utils.summarizer.Summarizer.summarize", new_callable=AsyncMock + ) as mock_summarize, + patch("openviking.utils.agfs_utils.create_agfs_client", return_value=mock_agfs), + patch("openviking.agfs_manager.AGFSManager.start"), + patch("openviking.agfs_manager.AGFSManager.stop"), + ): mock_summarize.return_value = {"status": "success"} - + client = AsyncOpenViking(path=str(test_config.parent)) await client.initialize() - + try: # 1. Test with build_index=True - await client.add_resource( - path=str(resource_file), - build_index=True, - wait=True - ) - + await client.add_resource(path=str(resource_file), build_index=True, wait=True) + # Verify summarizer called with skip_vectorization=False assert mock_summarize.call_count == 1 call_kwargs = mock_summarize.call_args.kwargs assert call_kwargs.get("skip_vectorization") is False - + mock_summarize.reset_mock() - + # 2. Test with build_index=False, summarize=True await client.add_resource( - path=str(resource_file), - build_index=False, - summarize=True, - wait=True + path=str(resource_file), build_index=False, summarize=True, wait=True ) - + # Verify summarizer called with skip_vectorization=True assert mock_summarize.call_count == 1 call_kwargs = mock_summarize.call_args.kwargs assert call_kwargs.get("skip_vectorization") is True - + mock_summarize.reset_mock() - + # 3. Test with build_index=False, summarize=False await client.add_resource( - path=str(resource_file), - build_index=False, - summarize=False, - wait=True + path=str(resource_file), build_index=False, summarize=False, wait=True ) - + # Verify summarizer NOT called mock_summarize.assert_not_called() - + finally: await client.close() OpenVikingConfigSingleton._instance = None diff --git a/tests/parse/test_ast_extractor.py b/tests/parse/test_ast_extractor.py index e7018c10..d887f7cf 100644 --- a/tests/parse/test_ast_extractor.py +++ b/tests/parse/test_ast_extractor.py @@ -10,36 +10,42 @@ # Helpers # --------------------------------------------------------------------------- + def _python_extractor(): from openviking.parse.parsers.code.ast.languages.python import PythonExtractor + return PythonExtractor() def _js_extractor(): from openviking.parse.parsers.code.ast.languages.js_ts import JsTsExtractor + return JsTsExtractor(lang="javascript") def _go_extractor(): from openviking.parse.parsers.code.ast.languages.go import GoExtractor + return GoExtractor() def _ts_extractor(): from openviking.parse.parsers.code.ast.languages.js_ts import JsTsExtractor + return JsTsExtractor(lang="typescript") def _csharp_extractor(): from openviking.parse.parsers.code.ast.languages.csharp import CSharpExtractor - return CSharpExtractor() + return CSharpExtractor() # --------------------------------------------------------------------------- # Python # --------------------------------------------------------------------------- + class TestPythonExtractor: SAMPLE = '''"""Module for parsing things. @@ -128,7 +134,9 @@ def test_multiline_params(self): # raw params may contain newlines, but to_text() must compact them assert "encoding" in methods["parse_async"].params text = sk.to_text() - assert "\n +" not in text.split("parse_async")[1].split("\n")[0] # no newline inside the signature line + assert ( + "\n +" not in text.split("parse_async")[1].split("\n")[0] + ) # no newline inside the signature line def test_top_level_function(self): sk = self.e.extract("test.py", self.SAMPLE) @@ -163,8 +171,9 @@ def test_to_text_verbose(self): # JavaScript # --------------------------------------------------------------------------- + class TestJavaScriptExtractor: - SAMPLE = ''' + SAMPLE = """ import React from "react"; import { useState, useEffect } from "react"; @@ -194,7 +203,7 @@ class Counter extends React.Component { function add(a, b) { return a + b; } -''' +""" def setup_method(self): self.e = _js_extractor() @@ -249,7 +258,7 @@ def test_to_text_verbose(self): assert "Maintains an internal count and exposes increment/decrement" in text def test_export_class(self): - code = ''' + code = """ /** Base utility class. * * Provides shared helper methods. @@ -258,7 +267,7 @@ def test_export_class(self): /** Log a message to the console. */ log(msg) { console.log(msg); } } -''' +""" sk = self.e.extract("utils.js", code) names = {c.name for c in sk.classes} assert "Utils" in names @@ -267,13 +276,13 @@ def test_export_class(self): assert any(m.name == "log" for m in cls.methods) def test_arrow_function(self): - code = ''' + code = """ /** Double a number. */ const double = (n) => n * 2; /** Negate a boolean. */ const negate = (b) => !b; -''' +""" sk = self.e.extract("math.js", code) names = {f.name for f in sk.functions} assert "double" in names @@ -286,8 +295,9 @@ def test_arrow_function(self): # Go # --------------------------------------------------------------------------- + class TestGoExtractor: - SAMPLE = ''' + SAMPLE = """ package main import ( @@ -312,7 +322,7 @@ class TestGoExtractor: fmt.Println("starting") return nil } -''' +""" def setup_method(self): self.e = _go_extractor() @@ -343,7 +353,9 @@ def test_method_receiver_not_params(self): def test_docstring_extracted(self): sk = self.e.extract("main.go", self.SAMPLE) fns = {f.name: f for f in sk.functions} - assert "NewServer creates a Server with the given host and port." in fns["NewServer"].docstring + assert ( + "NewServer creates a Server with the given host and port." in fns["NewServer"].docstring + ) assert "Returns a pointer to the initialized Server." in fns["NewServer"].docstring structs = {c.name: c for c in sk.classes} assert "Server" in structs @@ -369,13 +381,15 @@ def test_to_text_verbose(self): # Java # --------------------------------------------------------------------------- + def _java_extractor(): from openviking.parse.parsers.code.ast.languages.java import JavaExtractor + return JavaExtractor() class TestJavaExtractor: - SAMPLE = ''' + SAMPLE = """ import java.util.List; import java.util.Optional; @@ -408,7 +422,7 @@ class TestJavaExtractor: return a - b; } } -''' +""" def setup_method(self): self.e = _java_extractor() @@ -461,6 +475,7 @@ def test_to_text_verbose(self): # C# # --------------------------------------------------------------------------- + class TestCSharpExtractor: SAMPLE = """ using System; @@ -546,7 +561,7 @@ def test_to_text_verbose(self): assert "First operand" in text def test_file_scoped_namespace(self): - code = ''' + code = """ using System; namespace MyApp.Services; @@ -558,13 +573,13 @@ def test_file_scoped_namespace(self): return a + b; } } -''' +""" sk = self.e.extract("Calculator.cs", code) names = {c.name for c in sk.classes} assert "Calculator" in names def test_property_accessor_signature(self): - code = ''' + code = """ public class Calculator { /// @@ -572,7 +587,7 @@ def test_property_accessor_signature(self): /// public int Result { get; set; } } -''' +""" sk = self.e.extract("Calculator.cs", code) cls = next(c for c in sk.classes if c.name == "Calculator") methods = {m.name: m for m in cls.methods} @@ -585,13 +600,15 @@ def test_property_accessor_signature(self): # C/C++ # --------------------------------------------------------------------------- + def _cpp_extractor(): from openviking.parse.parsers.code.ast.languages.cpp import CppExtractor + return CppExtractor() class TestCppExtractor: - SAMPLE = ''' + SAMPLE = """ #include #include @@ -627,7 +644,7 @@ class Stack { int add(int a, int b) { return a + b; } -''' +""" def setup_method(self): self.e = _cpp_extractor() @@ -693,13 +710,15 @@ def test_to_text_verbose(self): # Rust # --------------------------------------------------------------------------- + def _rust_extractor(): from openviking.parse.parsers.code.ast.languages.rust import RustExtractor + return RustExtractor() class TestRustExtractor: - SAMPLE = ''' + SAMPLE = """ use std::collections::HashMap; use std::io::{self, Read}; @@ -736,7 +755,7 @@ class TestRustExtractor: pub fn factorial(n: u64) -> u64 { if n == 0 { 1 } else { n * factorial(n - 1) } } -''' +""" def setup_method(self): self.e = _rust_extractor() @@ -792,11 +811,11 @@ def test_to_text_verbose(self): assert "Panics if n is negative." in text - # --------------------------------------------------------------------------- # Skeleton.to_text() — verbose vs compact # --------------------------------------------------------------------------- + class TestSkeletonToText: MULTILINE_DOC = "First line summary.\n\nMore details here.\nArgs:\n x: an integer." @@ -808,9 +827,10 @@ def _make_skeleton(self): imports=["os", "sys"], classes=[ ClassSkeleton( - name="Foo", bases=["Base"], + name="Foo", + bases=["Base"], docstring=self.MULTILINE_DOC, - methods=[FunctionSig("run", "self", "None", self.MULTILINE_DOC)] + methods=[FunctionSig("run", "self", "None", self.MULTILINE_DOC)], ) ], functions=[FunctionSig("helper", "x: int", "bool", self.MULTILINE_DOC)], @@ -818,8 +838,12 @@ def _make_skeleton(self): def test_empty_skeleton(self): sk = CodeSkeleton( - file_name="empty.py", language="Python", - module_doc="", imports=[], classes=[], functions=[] + file_name="empty.py", + language="Python", + module_doc="", + imports=[], + classes=[], + functions=[], ) assert "# empty.py [Python]" in sk.to_text() @@ -844,7 +868,8 @@ def test_verbose_full_docstring(self): def test_verbose_single_line_doc_no_extra_quotes(self): sk = CodeSkeleton( - file_name="bar.py", language="Python", + file_name="bar.py", + language="Python", module_doc="Single line.", imports=[], classes=[ClassSkeleton("Bar", [], "One liner.", [])], @@ -862,8 +887,9 @@ def test_verbose_single_line_doc_no_extra_quotes(self): # TypeScript # --------------------------------------------------------------------------- + class TestTypeScriptExtractor: - SAMPLE = ''' + SAMPLE = """ import { Observable } from "rxjs"; import { HttpClient } from "@angular/common/http"; @@ -898,7 +924,7 @@ class TodoService { function validate(title: string): boolean { return title.length > 0 && title.length < 100; } -''' +""" def setup_method(self): self.e = _ts_extractor() @@ -965,9 +991,11 @@ def test_to_text_verbose(self): # ASTExtractor dispatch # --------------------------------------------------------------------------- + class TestASTExtractorDispatch: def setup_method(self): from openviking.parse.parsers.code.ast.extractor import ASTExtractor + self.extractor = ASTExtractor() def test_python_dispatch(self): @@ -977,13 +1005,13 @@ def test_python_dispatch(self): assert "def foo" in text def test_go_dispatch(self): - code = 'package main\n\n// Run starts the app.\nfunc Run() error {\n return nil\n}\n' + code = "package main\n\n// Run starts the app.\nfunc Run() error {\n return nil\n}\n" text = self.extractor.extract_skeleton("main.go", code) assert "# main.go [Go]" in text assert "Run" in text def test_csharp_dispatch(self): - code = 'namespace Demo;\n\npublic class Util { public int Add(int a, int b) { return a + b; } }\n' + code = "namespace Demo;\n\npublic class Util { public int Add(int a, int b) { return a + b; } }\n" text = self.extractor.extract_skeleton("util.cs", code) assert "# util.cs [C#]" in text assert "class Util" in text diff --git a/tests/server/test_api_filesystem.py b/tests/server/test_api_filesystem.py index 35ecd394..79058d37 100644 --- a/tests/server/test_api_filesystem.py +++ b/tests/server/test_api_filesystem.py @@ -7,9 +7,7 @@ async def test_ls_root(client: httpx.AsyncClient): - resp = await client.get( - "/api/v1/fs/ls", params={"uri": "viking://"} - ) + resp = await client.get("/api/v1/fs/ls", params={"uri": "viking://"}) assert resp.status_code == 200 body = resp.json() assert body["status"] == "ok" @@ -62,9 +60,7 @@ async def test_mkdir_and_ls(client: httpx.AsyncClient): async def test_tree(client: httpx.AsyncClient): - resp = await client.get( - "/api/v1/fs/tree", params={"uri": "viking://"} - ) + resp = await client.get("/api/v1/fs/tree", params={"uri": "viking://"}) assert resp.status_code == 200 body = resp.json() assert body["status"] == "ok" @@ -72,9 +68,7 @@ async def test_tree(client: httpx.AsyncClient): async def test_stat_after_add_resource(client_with_resource): client, uri = client_with_resource - resp = await client.get( - "/api/v1/fs/stat", params={"uri": uri} - ) + resp = await client.get("/api/v1/fs/stat", params={"uri": uri}) assert resp.status_code == 200 body = resp.json() assert body["status"] == "ok" @@ -92,9 +86,7 @@ async def test_stat_not_found(client: httpx.AsyncClient): async def test_rm_resource(client_with_resource): client, uri = client_with_resource - resp = await client.request( - "DELETE", "/api/v1/fs", params={"uri": uri, "recursive": True} - ) + resp = await client.request("DELETE", "/api/v1/fs", params={"uri": uri, "recursive": True}) assert resp.status_code == 200 assert resp.json()["status"] == "ok" diff --git a/tests/server/test_api_relations.py b/tests/server/test_api_relations.py index 02fd3324..cad2efda 100644 --- a/tests/server/test_api_relations.py +++ b/tests/server/test_api_relations.py @@ -6,9 +6,7 @@ async def test_get_relations_empty(client_with_resource): client, uri = client_with_resource - resp = await client.get( - "/api/v1/relations", params={"uri": uri} - ) + resp = await client.get("/api/v1/relations", params={"uri": uri}) assert resp.status_code == 200 body = resp.json() assert body["status"] == "ok" @@ -41,9 +39,7 @@ async def test_link_and_get_relations(client_with_resource): assert resp.json()["status"] == "ok" # Verify link exists - resp = await client.get( - "/api/v1/relations", params={"uri": uri} - ) + resp = await client.get("/api/v1/relations", params={"uri": uri}) assert resp.status_code == 200 body = resp.json() assert body["status"] == "ok" diff --git a/tests/server/test_api_resources.py b/tests/server/test_api_resources.py index bd6f9e0f..013c6baa 100644 --- a/tests/server/test_api_resources.py +++ b/tests/server/test_api_resources.py @@ -8,9 +8,7 @@ from tests.server.conftest import SAMPLE_MD_CONTENT -async def test_add_resource_success( - client: httpx.AsyncClient, sample_markdown_file -): +async def test_add_resource_success(client: httpx.AsyncClient, sample_markdown_file): resp = await client.post( "/api/v1/resources", json={ @@ -26,9 +24,7 @@ async def test_add_resource_success( assert body["result"]["root_uri"].startswith("viking://") -async def test_add_resource_with_wait( - client: httpx.AsyncClient, sample_markdown_file -): +async def test_add_resource_with_wait(client: httpx.AsyncClient, sample_markdown_file): resp = await client.post( "/api/v1/resources", json={ @@ -54,9 +50,7 @@ async def test_add_resource_file_not_found(client: httpx.AsyncClient): assert "errors" in body["result"] and len(body["result"]["errors"]) > 0 -async def test_add_resource_with_target( - client: httpx.AsyncClient, sample_markdown_file -): +async def test_add_resource_with_target(client: httpx.AsyncClient, sample_markdown_file): resp = await client.post( "/api/v1/resources", json={ @@ -81,9 +75,7 @@ async def test_wait_processed_empty_queue(client: httpx.AsyncClient): assert body["status"] == "ok" -async def test_wait_processed_after_add( - client: httpx.AsyncClient, sample_markdown_file -): +async def test_wait_processed_after_add(client: httpx.AsyncClient, sample_markdown_file): await client.post( "/api/v1/resources", json={"path": str(sample_markdown_file), "reason": "test"}, diff --git a/tests/server/test_api_search.py b/tests/server/test_api_search.py index 64ad6c1a..834c0437 100644 --- a/tests/server/test_api_search.py +++ b/tests/server/test_api_search.py @@ -66,9 +66,7 @@ async def test_search_basic(client_with_resource): async def test_search_with_session(client_with_resource): client, uri = client_with_resource # Create a session first - sess_resp = await client.post( - "/api/v1/sessions", json={"user": "test"} - ) + sess_resp = await client.post("/api/v1/sessions", json={"user": "test"}) session_id = sess_resp.json()["result"]["session_id"] resp = await client.post( diff --git a/tests/server/test_server_health.py b/tests/server/test_server_health.py index 5f66ea7e..fff93fe2 100644 --- a/tests/server/test_server_health.py +++ b/tests/server/test_server_health.py @@ -30,9 +30,7 @@ async def test_process_time_header(client: httpx.AsyncClient): async def test_openviking_error_handler(client: httpx.AsyncClient): """Requesting a non-existent resource should return structured error.""" - resp = await client.get( - "/api/v1/fs/stat", params={"uri": "viking://nonexistent/path"} - ) + resp = await client.get("/api/v1/fs/stat", params={"uri": "viking://nonexistent/path"}) assert resp.status_code in (404, 500) body = resp.json() assert body["status"] == "error" diff --git a/tests/storage/mock_backend.py b/tests/storage/mock_backend.py index 00d565ca..38737a1c 100644 --- a/tests/storage/mock_backend.py +++ b/tests/storage/mock_backend.py @@ -6,11 +6,13 @@ from openviking.storage.vectordb_adapters.base import CollectionAdapter from openviking.storage.vectordb.collection.collection import Collection + class MockCollectionAdapter(CollectionAdapter): """ Mock adapter for testing dynamic loading. Inherits from CollectionAdapter and wraps MockCollection. """ + def __init__(self, collection_name: str, custom_param1: str = "", custom_param2: int = 0): super().__init__(collection_name=collection_name) self.mode = "mock" @@ -23,32 +25,41 @@ def from_config(cls, config: Any) -> "MockCollectionAdapter": return cls( collection_name=config.name or "mock_collection", custom_param1=custom_params.get("custom_param1", ""), - custom_param2=custom_params.get("custom_param2", 0) + custom_param2=custom_params.get("custom_param2", 0), ) def _load_existing_collection_if_needed(self) -> None: if self._collection is None: - # Create a dummy collection wrapping MockCollection - self._collection = MockCollection(self.custom_param1, self.custom_param2) + # Create a dummy collection wrapping MockCollection + self._collection = MockCollection(self.custom_param1, self.custom_param2) def _create_backend_collection(self, meta: Dict[str, Any]) -> Collection: return MockCollection(self.custom_param1, self.custom_param2) + class MockCollection(ICollection): - def __init__(self, custom_param1: str, custom_param2: int, meta_data: Optional[Dict[str, Any]] = None, **kwargs): + def __init__( + self, + custom_param1: str, + custom_param2: int, + meta_data: Optional[Dict[str, Any]] = None, + **kwargs, + ): super().__init__() self.meta_data = meta_data if meta_data is not None else {} - + self.custom_param1 = custom_param1 self.custom_param2 = custom_param2 - + # Store extra kwargs (including host/headers if passed but not used explicitly) self.kwargs = kwargs - + # Verify that we can access values passed during initialization if self.meta_data and "test_verification" in self.meta_data: - print(f"MockCollection initialized with custom_param1={self.custom_param1}, custom_param2={self.custom_param2}, kwargs={kwargs}") - + print( + f"MockCollection initialized with custom_param1={self.custom_param1}, custom_param2={self.custom_param2}, kwargs={kwargs}" + ) + def update(self, fields: Optional[Dict[str, Any]] = None, description: Optional[str] = None): raise NotImplementedError("MockCollection.update is not supported") @@ -179,4 +190,4 @@ def aggregate_data( filters: Optional[Dict[str, Any]] = None, cond: Optional[Dict[str, Any]] = None, ) -> AggregateResult: - raise NotImplementedError("MockCollection.aggregate_data is not supported") \ No newline at end of file + raise NotImplementedError("MockCollection.aggregate_data is not supported") diff --git a/tests/storage/test_vectordb_adaptor.py b/tests/storage/test_vectordb_adaptor.py index 22146b41..5313f98d 100644 --- a/tests/storage/test_vectordb_adaptor.py +++ b/tests/storage/test_vectordb_adaptor.py @@ -15,21 +15,19 @@ import shutil import tempfile + class TestAdapterLoading(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() self.config_path = os.path.join(self.test_dir, "ov.conf") - + # Create a valid config file config_data = { "storage": { "vectordb": { "backend": "tests.storage.mock_backend.MockCollectionAdapter", "name": "mock_test_collection", - "custom_params": { - "custom_param1": "val1", - "custom_param2": 123 - } + "custom_params": {"custom_param1": "val1", "custom_param2": 123}, } }, "embedding": { @@ -37,9 +35,9 @@ def setUp(self): "provider": "openai", "model": "text-embedding-3-small", "api_key": "mock-key", - "dimension": 1536 + "dimension": 1536, } - } + }, } with open(self.config_path, "w") as f: json.dump(config_data, f) @@ -57,7 +55,7 @@ def test_dynamic_loading_mock_adapter(self): """ # Load config from the temporary file OpenVikingConfigSingleton.initialize(config_path=self.config_path) - + config = get_openviking_config().storage.vectordb # Verify that custom params are loaded @@ -67,23 +65,25 @@ def test_dynamic_loading_mock_adapter(self): try: adapter = create_collection_adapter(config) - + self.assertEqual(adapter.__class__.__name__, "MockCollectionAdapter") self.assertEqual(adapter.mode, "mock") self.assertEqual(adapter.collection_name, "mock_test_collection") self.assertEqual(adapter.custom_param1, "val1") self.assertEqual(adapter.custom_param2, 123) - + # Verify internal behavior exists = adapter.collection_exists() self.assertTrue(exists) - + print("Successfully loaded MockCollectionAdapter dynamically from config file.") - + except Exception as e: import traceback + traceback.print_exc() self.fail(f"Failed to load adapter dynamically: {e}") + if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/storage/test_vectordb_collection_loading.py b/tests/storage/test_vectordb_collection_loading.py index 3c1770e1..c5c9bd9e 100644 --- a/tests/storage/test_vectordb_collection_loading.py +++ b/tests/storage/test_vectordb_collection_loading.py @@ -5,9 +5,13 @@ # Add open_test path to ensure modules can be found sys.path.insert(0, "/cloudide/workspace/open_test") -from openviking.storage.vectordb.project.vikingdb_project import get_or_create_vikingdb_project, VikingDBProject +from openviking.storage.vectordb.project.vikingdb_project import ( + get_or_create_vikingdb_project, + VikingDBProject, +) from openviking.storage.vectordb.collection.vikingdb_collection import VikingDBCollection + class TestDynamicLoading(unittest.TestCase): def test_default_loading(self): # Test with default configuration @@ -19,76 +23,71 @@ def test_default_loading(self): def test_explicit_loading(self): # Test with explicit configuration pointing to MockJoiner # MockJoiner is now in tests.storage.mock_backend - + # We assume tests package structure is available from /cloudide/workspace/open_test - + config = { - "Host": "test_host", + "Host": "test_host", "Headers": {"Auth": "Token"}, "CollectionClass": "tests.storage.mock_backend.MockCollection", - "CollectionArgs": { - "custom_param1": "custom_val", - "custom_param2": 123 - } + "CollectionArgs": {"custom_param1": "custom_val", "custom_param2": 123}, } project = get_or_create_vikingdb_project(config=config) - + from tests.storage.mock_backend import MockCollection + self.assertEqual(project.CollectionClass, MockCollection) self.assertEqual(project.host, "test_host") self.assertEqual(project.headers, {"Auth": "Token"}) - self.assertEqual(project.collection_args, {"custom_param1": "custom_val", "custom_param2": 123}) - + self.assertEqual( + project.collection_args, {"custom_param1": "custom_val", "custom_param2": 123} + ) + # Test collection creation to verify params are passed collection_name = "test_collection" meta_data = { "test_verification": True, "Host": "metadata_host", - "Headers": {"Meta": "Header"} + "Headers": {"Meta": "Header"}, } - + # The project wrapper will pass host, headers, meta_data, AND collection_args - kwargs = { - "host": project.host, - "headers": project.headers, - "meta_data": meta_data - } + kwargs = {"host": project.host, "headers": project.headers, "meta_data": meta_data} kwargs.update(project.collection_args) - + collection_instance = project.CollectionClass(**kwargs) - + # Verify custom params are set correctly self.assertEqual(collection_instance.custom_param1, "custom_val") self.assertEqual(collection_instance.custom_param2, 123) - + # Verify host/headers are in kwargs (since init doesn't take them explicitly anymore) self.assertEqual(collection_instance.kwargs.get("host"), "test_host") self.assertEqual(collection_instance.kwargs.get("headers"), {"Auth": "Token"}) - + print("Explicit loading test passed (MockCollection with custom params)") def test_kwargs_loading(self): # Test with CollectionArgs config = { - "Host": "test_host", + "Host": "test_host", "CollectionClass": "tests.storage.mock_backend.MockCollection", - "CollectionArgs": { - "custom_param1": "extra_value", - "custom_param2": 456 - } + "CollectionArgs": {"custom_param1": "extra_value", "custom_param2": 456}, } project = get_or_create_vikingdb_project(config=config) - - self.assertEqual(project.collection_args, {"custom_param1": "extra_value", "custom_param2": 456}) - + + self.assertEqual( + project.collection_args, {"custom_param1": "extra_value", "custom_param2": 456} + ) + # Manually verify instantiation with kwargs kwargs = { "host": project.host, "headers": project.headers, - "meta_data": {"test_verification": True} + "meta_data": {"test_verification": True}, } kwargs.update(project.collection_args) - + collection_instance = project.CollectionClass(**kwargs) self.assertEqual(collection_instance.custom_param1, "extra_value") self.assertEqual(collection_instance.custom_param2, 456) @@ -96,13 +95,11 @@ def test_kwargs_loading(self): def test_invalid_loading(self): # Test with invalid class path - config = { - "Host": "test_host", - "CollectionClass": "non.existent.module.Class" - } + config = {"Host": "test_host", "CollectionClass": "non.existent.module.Class"} with self.assertRaises(ImportError): get_or_create_vikingdb_project(config=config) print("Invalid loading test passed") -if __name__ == '__main__': + +if __name__ == "__main__": unittest.main() diff --git a/tests/test_code_hosting_utils.py b/tests/test_code_hosting_utils.py index a1270457..41b60b52 100644 --- a/tests/test_code_hosting_utils.py +++ b/tests/test_code_hosting_utils.py @@ -34,7 +34,9 @@ def _mock_config(): sys.modules[_config_mod_name].get_openviking_config = _mock_config # type: ignore[attr-defined] # Load code_hosting_utils directly from file to avoid the heavy openviking/__init__.py chain -_module_path = Path(__file__).resolve().parents[1] / "openviking" / "utils" / "code_hosting_utils.py" +_module_path = ( + Path(__file__).resolve().parents[1] / "openviking" / "utils" / "code_hosting_utils.py" +) _spec = importlib.util.spec_from_file_location("openviking.utils.code_hosting_utils", _module_path) _module = importlib.util.module_from_spec(_spec) sys.modules["openviking.utils.code_hosting_utils"] = _module diff --git a/tests/utils/mock_agfs.py b/tests/utils/mock_agfs.py index af8fa649..24df4962 100644 --- a/tests/utils/mock_agfs.py +++ b/tests/utils/mock_agfs.py @@ -2,45 +2,49 @@ from pathlib import Path from unittest.mock import MagicMock + class MockLocalAGFS: """ A mock implementation of AGFSClient that operates on a local directory. Useful for tests where running a real AGFS server is not feasible or desired. """ + def __init__(self, config=None, root_path=None): self.config = config self.root = Path(root_path) if root_path else Path("/tmp/viking_data") self.root.mkdir(parents=True, exist_ok=True) - + def _resolve(self, path): if str(path).startswith("viking://"): path = str(path).replace("viking://", "") if str(path).startswith("/"): path = str(path)[1:] return self.root / path - + def exists(self, path, ctx=None): return self._resolve(path).exists() - + def mkdir(self, path, ctx=None, parents=True, exist_ok=True): self._resolve(path).mkdir(parents=parents, exist_ok=exist_ok) - + def ls(self, path, ctx=None, **kwargs): p = self._resolve(path) if not p.exists(): return [] res = [] for item in p.iterdir(): - res.append({ - "name": item.name, - "isDir": item.is_dir(), # Note: JS style camelCase for some APIs - "type": "directory" if item.is_dir() else "file", - "size": item.stat().st_size if item.is_file() else 0, - "mtime": item.stat().st_mtime, - "uri": f"viking://{path}/{item.name}".replace("//", "/") - }) + res.append( + { + "name": item.name, + "isDir": item.is_dir(), # Note: JS style camelCase for some APIs + "type": "directory" if item.is_dir() else "file", + "size": item.stat().st_size if item.is_file() else 0, + "mtime": item.stat().st_mtime, + "uri": f"viking://{path}/{item.name}".replace("//", "/"), + } + ) return res - + def writeto(self, path, content, ctx=None, **kwargs): p = self._resolve(path) p.parent.mkdir(parents=True, exist_ok=True) @@ -55,13 +59,13 @@ def write(self, path, content, ctx=None, **kwargs): def write_file(self, path, content, ctx=None, **kwargs): return self.writeto(path, content, ctx, **kwargs) - + def read_file(self, path, ctx=None, **kwargs): p = self._resolve(path) if not p.exists(): raise FileNotFoundError(path) return p.read_bytes() - + def read(self, path, ctx=None, **kwargs): return self.read_file(path, ctx, **kwargs) @@ -78,23 +82,19 @@ def rm(self, path, recursive=False, ctx=None): def delete_temp(self, path, ctx=None): self.rm(path, recursive=True, ctx=ctx) - + def mv(self, src, dst, ctx=None): s = self._resolve(src) d = self._resolve(dst) d.parent.mkdir(parents=True, exist_ok=True) shutil.move(str(s), str(d)) - + def stat(self, path, ctx=None): p = self._resolve(path) if not p.exists(): raise FileNotFoundError(path) s = p.stat() - return { - "size": s.st_size, - "mtime": s.st_mtime, - "is_dir": p.is_dir() - } - + return {"size": s.st_size, "mtime": s.st_mtime, "is_dir": p.is_dir()} + def bind_request_context(self, ctx): - return MagicMock(__enter__=lambda x: None, __exit__=lambda x,y,z: None) + return MagicMock(__enter__=lambda x: None, __exit__=lambda x, y, z: None) diff --git a/tests/vectordb/test_openviking_vectordb.py b/tests/vectordb/test_openviking_vectordb.py index 805950d3..0562a97b 100644 --- a/tests/vectordb/test_openviking_vectordb.py +++ b/tests/vectordb/test_openviking_vectordb.py @@ -288,9 +288,11 @@ def test_filters_update_delete_recall(self): }, ), self._expected_ids( - lambda item: item["context_type"] == "text" - and "tag_b" in item["tags"] - and item["is_leaf"] is False + lambda item: ( + item["context_type"] == "text" + and "tag_b" in item["tags"] + and item["is_leaf"] is False + ) ), ) self.assertEqual(