mirror of
https://github.com/GNS3/gns3-server.git
synced 2026-08-27 12:30:13 +03:00
## Summary Add a complete fault injection system for GNS3 Copilot, migrate all skills from local Python files to an external Git repository with hot reload support, and restructure Copilot API under /copilot/. ## Key Changes ### Fault Injection - New troubleshooting_injection mode with InjectionSkillsTool - 368 fault scenarios across 39 protocol categories - Context-based filtering (LLM must pass topology protocols) ### External Skills Repository - SkillsManager: Git clone/pull, version tracking, smart updates - SkillsLoader: YAML skills + Markdown prompts from external repo - Hot reload via POST /copilot/reload/skills - Configurable via gns3_server.conf ### Architecture - API unified under /copilot/ prefix - SkillsManager moved from Controller to agent module - Lazy initialization with startup background preload - Per-command Git timeout, smart update checks - Forbidden commands hot-reloadable from external repo - 32 INFO logs downgraded to DEBUG
783 lines
29 KiB
Python
783 lines
29 KiB
Python
# SPDX-License-Identifier: GPL-3.0-or-later
|
|
#
|
|
# GNS3-Copilot - AI-powered Network Lab Assistant for GNS3
|
|
#
|
|
# This file is part of GNS3-Copilot project.
|
|
#
|
|
# GNS3-Copilot is free software: you can redistribute it and/or modify it
|
|
# under the terms of the GNU General Public License as published by the
|
|
# Free Software Foundation, either version 3 of the License, or (at your
|
|
# option) any later version.
|
|
#
|
|
# GNS3-Copilot is distributed in the hope that it will be useful, but
|
|
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
|
|
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
|
|
# for more details.
|
|
#
|
|
# You should have received a copy of the GNU General Public License
|
|
# along with GNS3-Copilot. If not, see <https://www.gnu.org/licenses/>.
|
|
#
|
|
# Copyright (C) 2025 Yue Guobin (岳国宾)
|
|
# Author: Yue Guobin (岳国宾)
|
|
#
|
|
# Project Home: https://github.com/yueguobin/gns3-copilot
|
|
#
|
|
"""
|
|
|
|
GNS3 Copilot Agent Service
|
|
|
|
Provides project-level Agent instances with SQLite checkpoint management.
|
|
Each project has its own AgentService with a dedicated checkpoint database
|
|
in the project directory.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
from datetime import datetime
|
|
from typing import Any
|
|
from typing import AsyncGenerator
|
|
from typing import Dict
|
|
from typing import List
|
|
from typing import Optional
|
|
from uuid import uuid4
|
|
|
|
import aiosqlite
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
|
|
from gns3server.agent.gns3_copilot.agent.gns3_copilot import agent_builder
|
|
from gns3server.agent.gns3_copilot.chat_sessions_repository import (
|
|
ChatSessionsRepository,
|
|
)
|
|
from gns3server.agent.gns3_copilot.gns3_client.context_helpers import (
|
|
set_current_jwt_token,
|
|
)
|
|
from gns3server.agent.gns3_copilot.gns3_client.context_helpers import (
|
|
set_current_llm_config,
|
|
)
|
|
from gns3server.agent.gns3_copilot.utils.error_handler import format_error_message
|
|
from gns3server.agent.gns3_copilot.utils.message_converters import (
|
|
convert_langchain_to_openai,
|
|
)
|
|
from gns3server.agent.gns3_copilot.utils.tool_call_stream import (
|
|
ToolCallStreamAccumulator,
|
|
)
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class AgentService:
|
|
"""
|
|
Project-level Agent Service with async checkpoint management.
|
|
|
|
Manages a LangGraph agent instance with SQLite-based state persistence
|
|
for a single GNS3 project.
|
|
"""
|
|
|
|
def __init__(self, project_path: str):
|
|
"""
|
|
Initialize AgentService for a project.
|
|
|
|
Args:
|
|
project_path: Path to the GNS3 project directory
|
|
"""
|
|
self.project_path = project_path
|
|
self._checkpointer: Optional[AsyncSqliteSaver] = None
|
|
self._checkpointer_conn: Optional[aiosqlite.Connection] = None
|
|
self._checkpointer_path: Optional[str] = None
|
|
self._graph = None
|
|
self._init_lock = asyncio.Lock()
|
|
self._initialized = False
|
|
|
|
def abort_session(self, session_id: str):
|
|
"""
|
|
Signal abort for a session.
|
|
|
|
Args:
|
|
session_id: Session identifier
|
|
"""
|
|
from gns3server.agent.gns3_copilot.agent.gns3_copilot import (
|
|
set_abort_flag,
|
|
)
|
|
|
|
set_abort_flag(session_id)
|
|
log.info("Abort signal set for session: %s", session_id)
|
|
|
|
def _get_checkpoint_dir(self) -> str:
|
|
"""Get or create the checkpoint directory for this project."""
|
|
checkpoint_dir = os.path.join(self.project_path, "gns3-copilot")
|
|
os.makedirs(checkpoint_dir, exist_ok=True)
|
|
return checkpoint_dir
|
|
|
|
async def _get_checkpointer(self) -> AsyncSqliteSaver:
|
|
"""
|
|
Get or create the SQLite checkpointer for this project.
|
|
|
|
Returns:
|
|
AsyncSqliteSaver instance
|
|
"""
|
|
async with self._init_lock:
|
|
if self._checkpointer is not None:
|
|
return self._checkpointer
|
|
|
|
checkpoint_dir = self._get_checkpoint_dir()
|
|
checkpointer_path = os.path.join(
|
|
checkpoint_dir, "copilot_checkpoints.db"
|
|
)
|
|
|
|
log.debug("Creating checkpointer at: %s", checkpointer_path)
|
|
|
|
# Close existing connection if switching projects
|
|
if self._checkpointer_conn:
|
|
try:
|
|
await self._checkpointer_conn.close()
|
|
log.debug("Closed previous checkpointer connection")
|
|
except Exception as e:
|
|
log.warning(
|
|
"Error closing old checkpointer connection: %s", e
|
|
)
|
|
|
|
# Create new connection
|
|
conn = await aiosqlite.connect(checkpointer_path)
|
|
# Enable WAL mode for better concurrent performance
|
|
await conn.execute("PRAGMA journal_mode=WAL;")
|
|
self._checkpointer_conn = (
|
|
conn # Save connection reference to prevent GC
|
|
)
|
|
self._checkpointer = AsyncSqliteSaver(conn)
|
|
|
|
# CRITICAL: Initialize database schema
|
|
await self._checkpointer.setup()
|
|
|
|
# Create chat_sessions table in the same database
|
|
await self._create_chat_sessions_table(conn)
|
|
|
|
self._checkpointer_path = checkpointer_path
|
|
self._initialized = True
|
|
|
|
log.info("Project checkpointer created at: %s", checkpointer_path)
|
|
return self._checkpointer
|
|
|
|
async def _create_chat_sessions_table(self, conn: aiosqlite.Connection):
|
|
"""
|
|
Create the chat_sessions table in the checkpoint database.
|
|
|
|
Args:
|
|
conn: aiosqlite connection
|
|
"""
|
|
await conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS chat_sessions (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
thread_id TEXT UNIQUE NOT NULL,
|
|
user_id TEXT NOT NULL,
|
|
project_id TEXT NOT NULL,
|
|
title TEXT DEFAULT 'New Conversation',
|
|
|
|
-- Statistics
|
|
message_count INTEGER DEFAULT 0,
|
|
llm_calls_count INTEGER DEFAULT 0,
|
|
input_tokens INTEGER DEFAULT 0,
|
|
output_tokens INTEGER DEFAULT 0,
|
|
total_tokens INTEGER DEFAULT 0,
|
|
|
|
-- Timestamps
|
|
last_message_at TIMESTAMP,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
|
|
-- Reserved fields (JSON strings)
|
|
metadata TEXT DEFAULT '{}',
|
|
stats TEXT DEFAULT '{}',
|
|
|
|
-- Pin feature
|
|
pinned BOOLEAN DEFAULT FALSE
|
|
)
|
|
""")
|
|
|
|
# Create indexes
|
|
await conn.execute(
|
|
"CREATE INDEX IF NOT EXISTS idx_thread_id ON "
|
|
"chat_sessions(thread_id)"
|
|
)
|
|
await conn.execute(
|
|
"CREATE INDEX IF NOT EXISTS idx_user_project ON "
|
|
"chat_sessions(user_id, project_id)"
|
|
)
|
|
|
|
# Check if pinned column exists, add it if not (migration for existing
|
|
# databases)
|
|
cursor = await conn.execute("PRAGMA table_info(chat_sessions)")
|
|
columns = await cursor.fetchall()
|
|
column_names = [col[1] for col in columns]
|
|
|
|
if "pinned" not in column_names:
|
|
log.debug("Adding pinned column to existing chat_sessions table")
|
|
await conn.execute(
|
|
"ALTER TABLE chat_sessions ADD COLUMN pinned BOOLEAN DEFAULT "
|
|
"FALSE"
|
|
)
|
|
await conn.commit()
|
|
|
|
# Create pinned index (after column is guaranteed to exist)
|
|
await conn.execute(
|
|
"CREATE INDEX IF NOT EXISTS idx_pinned_updated ON "
|
|
"chat_sessions(pinned DESC, updated_at DESC)"
|
|
)
|
|
|
|
await conn.commit()
|
|
log.debug("chat_sessions table created in checkpoint database")
|
|
|
|
async def _get_graph(self):
|
|
"""Get or compile the LangGraph agent."""
|
|
if self._graph is None:
|
|
checkpointer = await self._get_checkpointer()
|
|
self._graph = agent_builder.compile(checkpointer=checkpointer)
|
|
log.info(
|
|
"LangGraph agent compiled for project: %s", self.project_path
|
|
)
|
|
return self._graph
|
|
|
|
async def stream_chat(
|
|
self,
|
|
message: str,
|
|
session_id: str,
|
|
project_id: Optional[str] = None,
|
|
user_id: Optional[str] = None,
|
|
jwt_token: Optional[str] = None,
|
|
mode: str = "text",
|
|
llm_config: Optional[Dict[str, Any]] = None,
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|
"""
|
|
Stream chat responses from the agent.
|
|
|
|
Args:
|
|
message: User message
|
|
session_id: Session/thread ID for conversation continuity
|
|
project_id: GNS3 project ID (optional, for context)
|
|
user_id: User ID for metadata tracking
|
|
jwt_token: JWT token for API authentication (optional)
|
|
mode: Interaction mode (default: "text")
|
|
llm_config: LLM configuration dict (provider, model, api_key,
|
|
etc.)
|
|
|
|
Yields:
|
|
Dict containing SSE-compatible response chunks
|
|
"""
|
|
log.info(
|
|
"Stream chat started: project_id=%s, user_id=%s, session_id=%s, "
|
|
"mode=%s",
|
|
project_id,
|
|
user_id,
|
|
session_id,
|
|
mode,
|
|
)
|
|
|
|
# Ensure checkpointer is initialized
|
|
if not self._checkpointer_conn:
|
|
await self._get_checkpointer()
|
|
|
|
# Get or create chat session
|
|
repo = ChatSessionsRepository(self._checkpointer_conn)
|
|
session = await repo.get_session_by_thread(session_id)
|
|
is_new_session = session is None
|
|
|
|
if is_new_session:
|
|
# Create new session
|
|
copilot_mode = llm_config.get("copilot_mode", "teaching_assistant").lower()
|
|
session = await repo.create_session(
|
|
thread_id=session_id,
|
|
user_id=user_id or "",
|
|
project_id=project_id or "",
|
|
title="New Conversation",
|
|
copilot_mode=copilot_mode,
|
|
)
|
|
log.debug(
|
|
"Created new chat session: thread_id=%s, copilot_mode=%s",
|
|
session_id,
|
|
copilot_mode
|
|
)
|
|
|
|
# Set request-scoped context variables (memory-only, not persisted)
|
|
if jwt_token:
|
|
set_current_jwt_token(jwt_token)
|
|
log.debug("JWT token set in context")
|
|
if llm_config:
|
|
set_current_llm_config(llm_config)
|
|
log.debug(
|
|
"LLM config set in context: provider=%s, model=%s",
|
|
llm_config.get("provider"),
|
|
llm_config.get("model"),
|
|
)
|
|
|
|
# Build config - only thread-safe identifiers
|
|
# Determine recursion_limit based on copilot_mode
|
|
copilot_mode = llm_config.get("copilot_mode", "teaching_assistant").lower()
|
|
if copilot_mode == "troubleshooting_injection":
|
|
recursion_limit = 100 # Need more recursion depth for fault injection workflow
|
|
log.debug("Using extended recursion_limit for troubleshooting_injection mode: 100")
|
|
else:
|
|
recursion_limit = 25 # Default recursion limit
|
|
|
|
config = {
|
|
"configurable": {
|
|
"thread_id": session_id,
|
|
"project_id": project_id,
|
|
},
|
|
"metadata": {
|
|
"user_id": user_id,
|
|
"copilot_mode": copilot_mode, # Store mode in checkpoint metadata
|
|
},
|
|
"recursion_limit": recursion_limit,
|
|
}
|
|
|
|
# Build inputs
|
|
# Determine remaining_steps based on copilot_mode
|
|
# Troubleshooting injection mode needs more steps for multi-tool workflow
|
|
if copilot_mode == "troubleshooting_injection":
|
|
remaining_steps = 50 # Need more steps for fault injection workflow
|
|
log.debug("Using extended remaining_steps for troubleshooting_injection mode: 50")
|
|
else:
|
|
remaining_steps = 20 # Default for other modes
|
|
|
|
inputs = {
|
|
"messages": [
|
|
HumanMessage(
|
|
content=message,
|
|
id=str(uuid4()),
|
|
metadata={"created_at": datetime.utcnow().isoformat()},
|
|
)
|
|
],
|
|
"llm_calls": 0,
|
|
"remaining_steps": remaining_steps,
|
|
"session_id": session_id,
|
|
"abort": False,
|
|
}
|
|
|
|
# Clear abort flag for this session at the start of stream
|
|
from gns3server.agent.gns3_copilot.agent.gns3_copilot import (
|
|
clear_abort_flag,
|
|
)
|
|
|
|
clear_abort_flag(session_id)
|
|
log.debug("Abort flag cleared for session: %s", session_id)
|
|
|
|
# Get the compiled graph
|
|
graph = await self._get_graph()
|
|
log.debug("LangGraph graph obtained, starting stream")
|
|
|
|
# Track statistics for session update
|
|
message_count = 1 # User message
|
|
llm_calls_count = 0
|
|
input_tokens = 0
|
|
output_tokens = 0
|
|
last_message_at = datetime.utcnow().isoformat()
|
|
|
|
# Track if we've counted the AI response for this turn
|
|
ai_response_counted = False
|
|
tool_messages_counted = 0
|
|
|
|
# Initialize tool call stream accumulator for handling progressive
|
|
# tool call arguments
|
|
tool_call_accumulator = ToolCallStreamAccumulator()
|
|
|
|
# Track if stream was aborted
|
|
stream_aborted = False
|
|
|
|
# Stream events
|
|
try:
|
|
async for event in graph.astream_events(
|
|
inputs, config=config, version="v2"
|
|
):
|
|
event_type = event.get("event", "")
|
|
data = event.get("data", {})
|
|
|
|
# Track LLM calls and tokens
|
|
if event_type == "on_chat_model_start":
|
|
# Filter out title_generator_node from statistics
|
|
langgraph_node = event.get("metadata", {}).get(
|
|
"langgraph_node", ""
|
|
)
|
|
if langgraph_node != "title_generator_node":
|
|
llm_calls_count += 1
|
|
log.debug(
|
|
"LLM call started, count=%d", llm_calls_count
|
|
)
|
|
else:
|
|
log.debug(
|
|
"Skipping LLM call count for internal node: "
|
|
"title_generator_node"
|
|
)
|
|
|
|
elif event_type == "on_chat_model_end":
|
|
# Filter out title_generator_node from token counting
|
|
langgraph_node = event.get("metadata", {}).get(
|
|
"langgraph_node", ""
|
|
)
|
|
if langgraph_node == "title_generator_node":
|
|
log.debug(
|
|
"Skipping token counting for internal node: "
|
|
"title_generator_node"
|
|
)
|
|
else:
|
|
# Extract token usage from response metadata
|
|
# Try multiple possible locations where token usage
|
|
# might be stored
|
|
token_info_found = False
|
|
|
|
# Method 1: response.usage_metadata
|
|
response = data.get("response", {})
|
|
if hasattr(response, "usage_metadata"):
|
|
usage = response.usage_metadata
|
|
if usage:
|
|
input_tokens += usage.get("input_tokens", 0)
|
|
output_tokens += usage.get("output_tokens", 0)
|
|
token_info_found = True
|
|
|
|
# Method 2: output.usage_metadata
|
|
if not token_info_found:
|
|
output_msg = data.get("output", {})
|
|
if hasattr(output_msg, "usage_metadata"):
|
|
usage = output_msg.usage_metadata
|
|
if usage:
|
|
input_tokens += usage.get(
|
|
"input_tokens", 0
|
|
)
|
|
output_tokens += usage.get(
|
|
"output_tokens", 0
|
|
)
|
|
token_info_found = True
|
|
|
|
# Method 3: Check data directly for token usage fields
|
|
if not token_info_found:
|
|
if "input_tokens" in data:
|
|
input_tokens += data.get("input_tokens", 0)
|
|
if "output_tokens" in data:
|
|
output_tokens += data.get("output_tokens", 0)
|
|
if (
|
|
"input_tokens" in data
|
|
or "output_tokens" in data
|
|
):
|
|
token_info_found = True
|
|
|
|
# Count AI response as one message (only once per turn)
|
|
if not ai_response_counted:
|
|
message_count += 1
|
|
ai_response_counted = True
|
|
|
|
# Track tool messages
|
|
elif event_type == "on_tool_end":
|
|
message_count += 1 # Tool result message
|
|
tool_messages_counted += 1
|
|
log.debug(
|
|
"Tool message counted, message_count=%d", message_count
|
|
)
|
|
|
|
# Convert event to chunk for SSE streaming
|
|
# Use accumulator for on_chat_model_stream events to handle
|
|
# progressive tool calls
|
|
|
|
# Filter out internal nodes (title_generator_node) from
|
|
# streaming to frontend
|
|
langgraph_node = event.get("metadata", {}).get(
|
|
"langgraph_node", ""
|
|
)
|
|
if langgraph_node == "title_generator_node":
|
|
# Skip all events from the title_generator_node (internal
|
|
# use only)
|
|
log.debug(
|
|
"Skipping event from internal node: "
|
|
"title_generator_node"
|
|
)
|
|
continue
|
|
|
|
if event_type == "on_chat_model_stream":
|
|
chunks = tool_call_accumulator.process_event(event)
|
|
for chunk in chunks:
|
|
# Add session_id to each chunk
|
|
chunk["session_id"] = session_id
|
|
log.debug(
|
|
"Yielding accumulated chunk: type=%s",
|
|
chunk.get("type"),
|
|
)
|
|
yield chunk
|
|
else:
|
|
# Use stateless converter for other events
|
|
chunk = self._convert_event_to_chunk(event, session_id)
|
|
if chunk:
|
|
log.debug("Yielding chunk: type=%s", chunk.get("type"))
|
|
yield chunk
|
|
|
|
# Check if stream was aborted and yield tool_end events for aborted tools
|
|
from gns3server.agent.gns3_copilot.agent.gns3_copilot import (
|
|
check_abort_flag,
|
|
)
|
|
|
|
if check_abort_flag(session_id):
|
|
log.info("Stream aborted, yielding tool_end events: session_id=%s", session_id)
|
|
# Get final state to find aborted tool messages
|
|
final_state = await graph.aget_state(config)
|
|
if final_state and "messages" in final_state.values:
|
|
for msg in final_state.values["messages"]:
|
|
# Check if this is an aborted tool message
|
|
if hasattr(msg, "metadata") and msg.metadata.get("aborted"):
|
|
yield {
|
|
"type": "tool_end",
|
|
"tool_name": getattr(msg, "name", "unknown"),
|
|
"tool_output": msg.content if hasattr(msg, "content") else "",
|
|
"session_id": session_id,
|
|
}
|
|
|
|
# Update session statistics after successful stream
|
|
await repo.update_session(
|
|
thread_id=session_id,
|
|
message_count=message_count,
|
|
llm_calls_count=llm_calls_count,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
total_tokens=input_tokens + output_tokens,
|
|
last_message_at=last_message_at,
|
|
)
|
|
log.info(
|
|
"Session statistics updated: thread_id=%s, messages=%d, "
|
|
"llm_calls=%d, tokens=%d+%d=%d",
|
|
session_id,
|
|
message_count,
|
|
llm_calls_count,
|
|
input_tokens,
|
|
output_tokens,
|
|
input_tokens + output_tokens,
|
|
)
|
|
|
|
# Sync auto-generated title from checkpoint state
|
|
final_state = await graph.aget_state(config)
|
|
if final_state and "conversation_title" in final_state.values:
|
|
generated_title = final_state.values["conversation_title"]
|
|
current_session = await repo.get_session_by_thread(session_id)
|
|
if (
|
|
current_session
|
|
and current_session.title != generated_title
|
|
):
|
|
await repo.update_session(
|
|
thread_id=session_id, title=generated_title
|
|
)
|
|
log.info(
|
|
"Auto-generated title synced: thread_id=%s, title=%s",
|
|
session_id,
|
|
generated_title,
|
|
)
|
|
|
|
except Exception as e:
|
|
log.error("Error in stream_chat: %s", e, exc_info=True)
|
|
yield {
|
|
"type": "error",
|
|
"error": format_error_message(e),
|
|
"session_id": session_id,
|
|
}
|
|
|
|
def _convert_event_to_chunk(
|
|
self, event: Dict[str, Any], session_id: str
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Convert LangGraph event to API response chunk.
|
|
|
|
Args:
|
|
event: LangGraph event from astream_events
|
|
session_id: Session ID for the response
|
|
|
|
Returns:
|
|
Dict for SSE response or None if event should be filtered
|
|
|
|
Note:
|
|
on_chat_model_stream events are handled by
|
|
ToolCallStreamAccumulator before calling this method, so they are
|
|
not processed here.
|
|
"""
|
|
event_type = event.get("event", "")
|
|
data = event.get("data", {})
|
|
|
|
if event_type == "on_tool_start":
|
|
# Tool execution started
|
|
# Extract tool_call_id from event metadata to associate with
|
|
# tool_call event
|
|
tool_call_id = event.get("metadata", {}).get("tool_call_id", "")
|
|
return {
|
|
"type": "tool_start",
|
|
"tool_name": event.get("name", ""),
|
|
"tool_call_id": tool_call_id,
|
|
"session_id": session_id,
|
|
}
|
|
|
|
elif event_type == "on_tool_end":
|
|
# Tool execution completed
|
|
# Extract tool output and convert to JSON string
|
|
output = data.get("output", "")
|
|
# Convert output to JSON string if it's not already a string
|
|
# This ensures dict/list outputs are properly serialized for
|
|
# frontend parsing
|
|
if not isinstance(output, str):
|
|
output = json.dumps(output, ensure_ascii=False, indent=2)
|
|
return {
|
|
"type": "tool_end",
|
|
"tool_name": event.get("name", ""),
|
|
"tool_output": output,
|
|
"session_id": session_id,
|
|
}
|
|
|
|
return None
|
|
|
|
async def get_history(
|
|
self, session_id: str, limit: int = 100
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Get conversation history for a session.
|
|
|
|
Args:
|
|
session_id: Session/thread ID
|
|
limit: Maximum number of messages to retrieve
|
|
|
|
Returns:
|
|
Dict containing thread_id, title, and messages
|
|
"""
|
|
config = {"configurable": {"thread_id": session_id}}
|
|
|
|
try:
|
|
graph = await self._get_graph()
|
|
state = await graph.aget_state(config)
|
|
|
|
if state and "messages" in state.values:
|
|
messages = []
|
|
for msg in state.values["messages"][-limit:]:
|
|
messages.append(self._convert_message_to_dict(msg))
|
|
|
|
title = state.values.get(
|
|
"conversation_title", "New Conversation"
|
|
)
|
|
|
|
return {
|
|
"thread_id": session_id,
|
|
"title": title,
|
|
"messages": messages,
|
|
}
|
|
except Exception as e:
|
|
log.error("Error getting history: %s", e, exc_info=True)
|
|
|
|
return {
|
|
"thread_id": session_id,
|
|
"title": "New Conversation",
|
|
"messages": [],
|
|
}
|
|
|
|
def _convert_message_to_dict(self, msg) -> Dict[str, Any]:
|
|
"""Convert a LangChain message to OpenAI-compatible dict format."""
|
|
return convert_langchain_to_openai(msg)
|
|
|
|
async def list_sessions(
|
|
self, user_id: Optional[str] = None, copilot_mode: Optional[str] = None, limit: int = 100
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
List chat sessions for this project.
|
|
|
|
Args:
|
|
user_id: Filter by user ID (optional)
|
|
copilot_mode: Filter by copilot mode (optional)
|
|
limit: Maximum number of sessions to return
|
|
|
|
Returns:
|
|
List of session dictionaries
|
|
"""
|
|
if not self._checkpointer_conn:
|
|
await self._get_checkpointer()
|
|
|
|
repo = ChatSessionsRepository(self._checkpointer_conn)
|
|
sessions = await repo.list_sessions(user_id=user_id, copilot_mode=copilot_mode, limit=limit)
|
|
return [s.to_dict() for s in sessions]
|
|
|
|
async def delete_session(self, session_id: str) -> bool:
|
|
"""
|
|
Delete a chat session and its checkpoints.
|
|
|
|
Args:
|
|
session_id: Thread ID to delete
|
|
|
|
Returns:
|
|
True if deleted, False if not found
|
|
"""
|
|
if not self._checkpointer_conn:
|
|
await self._get_checkpointer()
|
|
|
|
repo = ChatSessionsRepository(self._checkpointer_conn)
|
|
return await repo.delete_session(session_id)
|
|
|
|
async def rename_session(
|
|
self, session_id: str, new_title: str
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Rename a chat session.
|
|
|
|
Args:
|
|
session_id: Thread ID
|
|
new_title: New title
|
|
|
|
Returns:
|
|
Updated session dictionary or None
|
|
"""
|
|
if not self._checkpointer_conn:
|
|
await self._get_checkpointer()
|
|
|
|
repo = ChatSessionsRepository(self._checkpointer_conn)
|
|
session = await repo.update_session(
|
|
thread_id=session_id, title=new_title
|
|
)
|
|
return session.to_dict() if session else None
|
|
|
|
async def pin_session(
|
|
self, session_id: str, pinned: bool = True
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Pin or unpin a chat session.
|
|
|
|
Args:
|
|
session_id: Thread ID
|
|
pinned: True to pin, False to unpin
|
|
|
|
Returns:
|
|
Updated session dictionary or None
|
|
"""
|
|
if not self._checkpointer_conn:
|
|
await self._get_checkpointer()
|
|
|
|
repo = ChatSessionsRepository(self._checkpointer_conn)
|
|
session = await repo.pin_session(thread_id=session_id, pinned=pinned)
|
|
return session.to_dict() if session else None
|
|
|
|
async def close(self):
|
|
"""
|
|
Close the checkpointer connection and cleanup resources.
|
|
"""
|
|
async with self._init_lock:
|
|
if self._checkpointer_conn:
|
|
try:
|
|
# Add timeout to prevent blocking on database close
|
|
await asyncio.wait_for(
|
|
self._checkpointer_conn.close(),
|
|
timeout=5.0
|
|
)
|
|
log.debug(
|
|
"Checkpointer connection closed for: %s",
|
|
self.project_path,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
log.warning(
|
|
"Checkpointer connection close timeout for: %s (forcing cleanup)",
|
|
self.project_path,
|
|
)
|
|
except Exception as e:
|
|
log.warning("Error closing checkpointer connection: %s", e)
|
|
finally:
|
|
self._checkpointer_conn = None
|
|
self._checkpointer = None
|
|
self._graph = None
|
|
self._initialized = False
|