chore(core-api): remove AI code - moved to core-ai service
Removed all AI/LLM functionality from core-api as it has been migrated to the dedicated core-ai service. Deleted: - src/controllers/ai_controller.py (chat completions, models, conversations) - src/agent/ (orchestrator, tools, prompts, streaming) - src/memory/ (manager, qdrant, buffer, schemas) - src/api/v1/ (chat, conversations, models, schemas) - tests/test_memory_*.py (3 test files) Removed dependencies: - google-adk, litellm, google-cloud-aiplatform - qdrant-client Kept: - tools_controller.py (web scraper for core-ai REST calls) - infrastructure_controller.py - health_controller.py - static_controller.py core-api is now purely for infrastructure management. All AI operations are handled by core-ai service. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -25,13 +25,5 @@ PyJWT[crypto]~=2.9.0
|
||||
python-jose[cryptography]~=3.3.0
|
||||
cryptography~=43.0.0
|
||||
|
||||
# Memory & Embeddings (using Ollama for embeddings - no local models needed)
|
||||
qdrant-client~=1.11.0
|
||||
|
||||
# Agent Framework - Google ADK (November 2025)
|
||||
# Using ADK with LiteLLM for Ollama compatibility
|
||||
# Pin dependencies to avoid slow resolution
|
||||
google-genai==1.17.0
|
||||
google-cloud-aiplatform[agent-engines]==1.95.1
|
||||
google-adk==1.3.0
|
||||
litellm==1.80.5
|
||||
# Note: AI/LLM functionality has been moved to core-ai service
|
||||
# core-api is now purely for infrastructure management
|
||||
|
||||
@@ -1,27 +0,0 @@
|
||||
"""
|
||||
Unified Agent Module
|
||||
|
||||
This module provides an intelligent agent that can handle infrastructure management,
|
||||
web search, and multi-step reasoning with transparent streaming output.
|
||||
"""
|
||||
|
||||
# Check if agent dependencies are available
|
||||
try:
|
||||
from .orchestrator import UnifiedAgent, get_unified_agent
|
||||
from .tools import get_agent_tools, ALL_TOOLS
|
||||
AGENT_AVAILABLE = True
|
||||
except ImportError as e:
|
||||
# ADK or dependencies not available
|
||||
AGENT_AVAILABLE = False
|
||||
UnifiedAgent = None
|
||||
get_unified_agent = None
|
||||
get_agent_tools = None
|
||||
ALL_TOOLS = []
|
||||
|
||||
__all__ = [
|
||||
"UnifiedAgent",
|
||||
"get_unified_agent",
|
||||
"get_agent_tools",
|
||||
"ALL_TOOLS",
|
||||
"AGENT_AVAILABLE",
|
||||
]
|
||||
@@ -1,259 +0,0 @@
|
||||
"""
|
||||
Agent Orchestrator - Unified intelligent agent with streaming reasoning using Google ADK
|
||||
|
||||
This orchestrator uses Google's Agent Development Kit (ADK) to create an agent that can:
|
||||
- Use tools to answer infrastructure questions
|
||||
- Stream thinking/reasoning output
|
||||
- Handle multi-step tasks
|
||||
- Work with local Ollama models via LiteLLM
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
from typing import AsyncIterator, Dict, Any, List, Optional
|
||||
from functools import lru_cache
|
||||
|
||||
try:
|
||||
from google.adk.agents import Agent
|
||||
from google.adk.models.lite_llm import LiteLlm
|
||||
from google.adk import Runner
|
||||
from google.adk.sessions import InMemorySessionService
|
||||
from google.genai.types import Content, Part
|
||||
ADK_AVAILABLE = True
|
||||
except ImportError:
|
||||
ADK_AVAILABLE = False
|
||||
logger_temp = logging.getLogger(__name__)
|
||||
logger_temp.error("Google ADK not installed! Run: pip install google-adk litellm")
|
||||
|
||||
from src.config import get_settings
|
||||
from src.agent.tools import get_agent_tools
|
||||
from src.agent.prompts import get_prompt
|
||||
import uuid
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class UnifiedAgent:
|
||||
"""
|
||||
Unified intelligent agent that handles all tool routing and reasoning using Google ADK
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
if not ADK_AVAILABLE:
|
||||
raise ImportError("Google ADK is not installed. Please install: pip install google-adk")
|
||||
|
||||
# Enable verbose logging for LiteLLM to debug prompts
|
||||
import litellm
|
||||
litellm.set_verbose = True
|
||||
logger.info("LiteLLM verbose logging enabled.")
|
||||
|
||||
self.settings = get_settings()
|
||||
self.tools = get_agent_tools()
|
||||
|
||||
# Initialize LiteLLM for Ollama (format: "ollama/model_name")
|
||||
model_name = self.settings.agent_model
|
||||
litellm_model = f"ollama/{model_name}"
|
||||
|
||||
logger.info(f"Initializing ADK with LiteLLM model: {litellm_model}")
|
||||
logger.info(f"Ollama base URL from settings: {self.settings.ollama_base_url}")
|
||||
|
||||
self.llm = LiteLlm(
|
||||
model=litellm_model,
|
||||
api_base=self.settings.ollama_base_url,
|
||||
response_format={"type": "text"},
|
||||
force_json=False,
|
||||
temperature=0.1,
|
||||
)
|
||||
|
||||
# Store agents with different prompts for A/B testing
|
||||
self._agents = {}
|
||||
|
||||
# Create default agent
|
||||
self.agent = self._get_agent(self.settings.system_prompt_variant)
|
||||
|
||||
# Create session service and runner for executing the agent
|
||||
session_service = InMemorySessionService()
|
||||
self.runner = Runner(
|
||||
app_name="portainer-core-api",
|
||||
agent=self.agent,
|
||||
session_service=session_service
|
||||
)
|
||||
|
||||
logger.info(f"Initialized ADK Agent with {len(self.tools)} tools using prompt variant: {self.settings.system_prompt_variant}")
|
||||
|
||||
def _get_agent(self, prompt_variant: str) -> Agent:
|
||||
"""Get or create an agent with a specific prompt variant"""
|
||||
if prompt_variant not in self._agents:
|
||||
system_prompt = get_prompt(prompt_variant)
|
||||
|
||||
self._agents[prompt_variant] = Agent(
|
||||
model=self.llm,
|
||||
name="tatlock",
|
||||
description="British butler assistant for technical household matters",
|
||||
instruction=system_prompt,
|
||||
tools=self.tools,
|
||||
)
|
||||
logger.info(f"Created ADK agent with prompt variant: {prompt_variant}")
|
||||
|
||||
return self._agents[prompt_variant]
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
message: str,
|
||||
conversation_history: List[Dict[str, str]] = None,
|
||||
stream: bool = True,
|
||||
prompt_variant: Optional[str] = None
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Process a chat message with streaming reasoning output
|
||||
|
||||
Args:
|
||||
message: User's message
|
||||
conversation_history: Previous conversation turns (optional)
|
||||
stream: Whether to stream intermediate steps
|
||||
prompt_variant: Override system prompt variant for A/B testing (optional)
|
||||
|
||||
Yields:
|
||||
Dict with keys:
|
||||
- type: "thinking" | "tool_call" | "tool_result" | "content" | "error"
|
||||
- content: The actual content
|
||||
- tool: Tool name (if type is tool_call)
|
||||
- model: Model being used (optional)
|
||||
"""
|
||||
try:
|
||||
# Generate session ID (use conversation history to maintain session)
|
||||
# For now, create a new session per request (stateless)
|
||||
user_id = "default_user"
|
||||
session_id = str(uuid.uuid4())
|
||||
|
||||
# Create session
|
||||
await self.runner.session_service.create_session(
|
||||
app_name="portainer-core-api",
|
||||
user_id=user_id,
|
||||
session_id=session_id
|
||||
)
|
||||
|
||||
# Create Content object from message
|
||||
new_message = Content(
|
||||
parts=[Part(text=message)],
|
||||
role="user"
|
||||
)
|
||||
|
||||
logger.info(f"🚀 Starting ADK Runner with message 🧠: {message[:50]}...")
|
||||
|
||||
# Track if we've seen any content
|
||||
has_content = False
|
||||
|
||||
logger.info("About to start async iteration over runner.run_async()")
|
||||
|
||||
# Stream from ADK Runner
|
||||
async for event in self.runner.run_async(
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
new_message=new_message
|
||||
):
|
||||
event_type_name = type(event).__name__
|
||||
logger.info(f"ADK Event: {event_type_name}")
|
||||
|
||||
# SPECIAL HANDLING for the model hallucinating a 'response' tool call.
|
||||
if event_type_name == "ToolCallStart" and event.tool_name == "response":
|
||||
try:
|
||||
answer = event.tool_input.get("answer", "")
|
||||
if answer:
|
||||
logger.info("📢 Intercepted 'response' tool. Delivering final answer.")
|
||||
yield {"type": "content", "content": answer}
|
||||
has_content = True
|
||||
# Gracefully exit the loop as this is the final response.
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing special 'response' tool call: {e}")
|
||||
break
|
||||
|
||||
# Map ADK events to our format
|
||||
event_type = type(event).__name__
|
||||
|
||||
if event_type == "ToolCallStart":
|
||||
# Tool is being called
|
||||
tool_name = getattr(event, "tool_name", "unknown")
|
||||
logger.info(f"🔧 TOOL CALL START: {tool_name}")
|
||||
|
||||
yield {
|
||||
"type": "tool_call",
|
||||
"tool": tool_name,
|
||||
"content": f"Using tool: {tool_name}..."
|
||||
}
|
||||
|
||||
elif event_type == "ToolCallEnd":
|
||||
# Tool execution completed
|
||||
tool_name = getattr(event, "tool_name", "unknown")
|
||||
logger.info(f"✅ TOOL CALL END: {tool_name}")
|
||||
|
||||
yield {
|
||||
"type": "tool_result",
|
||||
"content": "Tool execution complete"
|
||||
}
|
||||
|
||||
elif event_type == "ContentDelta":
|
||||
# Stream content tokens
|
||||
content = getattr(event, "content", "")
|
||||
if content:
|
||||
has_content = True
|
||||
yield {
|
||||
"type": "content",
|
||||
"content": content
|
||||
}
|
||||
|
||||
elif event_type == "AgentThinking":
|
||||
# Agent reasoning (if available)
|
||||
thinking = getattr(event, "content", "")
|
||||
if thinking:
|
||||
yield {
|
||||
"type": "thinking",
|
||||
"content": thinking
|
||||
}
|
||||
|
||||
logger.info(f"✅ ADK agent stream completed (has_content={has_content})")
|
||||
|
||||
# If no content was yielded, provide a default message
|
||||
if not has_content:
|
||||
logger.warning("No content generated by agent")
|
||||
yield {
|
||||
"type": "content",
|
||||
"content": "I apologize, but I wasn't able to generate a response."
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in ADK agent chat: {e}", exc_info=True)
|
||||
yield {
|
||||
"type": "error",
|
||||
"content": f"Sorry, I encountered an error: {str(e)}"
|
||||
}
|
||||
|
||||
async def chat_completion(
|
||||
self,
|
||||
message: str,
|
||||
conversation_history: List[Dict[str, str]] = None,
|
||||
prompt_variant: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Get a non-streaming response (for backwards compatibility)
|
||||
|
||||
Args:
|
||||
message: User's message
|
||||
conversation_history: Previous conversation turns (optional)
|
||||
prompt_variant: Override system prompt variant for A/B testing (optional)
|
||||
|
||||
Returns:
|
||||
The final response content
|
||||
"""
|
||||
final_content = ""
|
||||
async for chunk in self.chat(message, conversation_history, stream=True, prompt_variant=prompt_variant):
|
||||
if chunk["type"] == "content":
|
||||
final_content += chunk["content"]
|
||||
|
||||
return final_content if final_content else "I couldn't generate a response."
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def get_unified_agent() -> UnifiedAgent:
|
||||
"""Get cached unified agent instance"""
|
||||
return UnifiedAgent()
|
||||
@@ -1,363 +0,0 @@
|
||||
"""
|
||||
System Prompt Variants for A/B Testing
|
||||
|
||||
Each prompt is tested for:
|
||||
- Tool calling accuracy (does it call tools when needed?)
|
||||
- Response naturalness (does it sound like Tatlock?)
|
||||
- Instruction following (does it avoid announcing methods?)
|
||||
"""
|
||||
|
||||
PROMPTS = {
|
||||
"v1_verbose": """You are Tatlock, a British butler who assists with both conversation and household technical matters.
|
||||
|
||||
Your manner:
|
||||
- Polite and proper, addressing users as "sir"
|
||||
- Understated dry wit, the occasional sly remark
|
||||
- Economy of words - concise unless elaboration is warranted
|
||||
- Never fawning or obsequious
|
||||
|
||||
CRITICAL INSTRUCTIONS:
|
||||
|
||||
1. You have NO knowledge of your own - you must ALWAYS use tools to gather current factual information
|
||||
2. When you use tools, do NOT explain what you're doing - just use them silently
|
||||
3. The user will see progress indicators automatically (like "🔍 Searching web...")
|
||||
4. After gathering data, respond naturally in character with the results
|
||||
|
||||
Tools available to you:
|
||||
- web_search: For any current information, news, or facts from the internet
|
||||
- web_scrape: To read specific web pages in detail
|
||||
- list_services: To check which Docker containers are running
|
||||
- get_service_details: To inspect a specific service's status
|
||||
- list_domains: To check configured domains and proxies
|
||||
- get_system_status: To check CPU, memory, disk usage
|
||||
- get_current_time: To get current date/time (essential for "today" queries)
|
||||
- read_documentation: To read project docs
|
||||
|
||||
Decision tree for responses:
|
||||
|
||||
When asked about the current time or date:
|
||||
→ ALWAYS use get_current_time tool first
|
||||
→ Then respond: "Sir, the time is [time details]."
|
||||
|
||||
When asked about current facts, news, stocks, weather, research topics:
|
||||
→ ALWAYS call get_current_time first to know today's date
|
||||
→ Use web_search with the current date context
|
||||
→ Then respond: "Sir, I have examined [topic]. It appears [findings]..."
|
||||
|
||||
When asked about home server status, services, domains:
|
||||
→ Use appropriate infrastructure tool (list_services, get_service_details, etc.)
|
||||
→ Then respond: "Sir, I have checked the systems. [findings]..."
|
||||
|
||||
When asked conversational questions (opinions, jokes, how are you):
|
||||
→ NO TOOLS - just respond naturally in character
|
||||
|
||||
IMPORTANT: Never say things like "I shall use web search" or "I will call the tool" - just use the tool silently and then speak naturally about what you found.
|
||||
|
||||
Example flow (what the user sees):
|
||||
User: "What stock is trending highest today?"
|
||||
[🕐 Checking time...]
|
||||
[🔍 Searching web...]
|
||||
Agent: "Sir, I have examined today's markets. It appears NVIDIA is performing rather well at $142, up 3.2%. The gaming company turned AI purveyor continues to paint pretty pictures, as it were."
|
||||
|
||||
Remember: A proper butler doesn't announce his methods. He simply delivers results with appropriate wit and decorum.""",
|
||||
|
||||
"v2_concise": """You are Tatlock, a British butler who assists with household technical matters.
|
||||
|
||||
Your manner: Polite and proper, addressing users as "sir". Understated dry wit. Concise unless elaboration is warranted. Never fawning.
|
||||
|
||||
CRITICAL: You have NO internal knowledge. You MUST use tools to gather ALL factual information.
|
||||
|
||||
Tool usage rules:
|
||||
- Time/date questions → use get_current_time
|
||||
- Web searches, news, stocks, weather → use get_current_time first, then web_search
|
||||
- Docker services/containers → use list_services or get_service_details
|
||||
- System resources → use get_system_status
|
||||
- Simple conversation → no tools needed
|
||||
|
||||
After using tools, respond naturally without mentioning what tools you used. The user sees progress indicators automatically.
|
||||
|
||||
Example:
|
||||
User: "What time is it?"
|
||||
[You call get_current_time tool silently]
|
||||
You: "Sir, it's 14:35 on Monday, November 24th."
|
||||
|
||||
Remember: Always use tools for facts. Never guess or use your own knowledge.""",
|
||||
|
||||
"v3_imperative": """You are Tatlock, a proper British butler assisting with technical household matters.
|
||||
|
||||
Character: Polite, dry wit, concise, addresses users as "sir".
|
||||
|
||||
CRITICAL RULES:
|
||||
|
||||
1. You have NO knowledge of current time, dates, weather, news, or system status
|
||||
2. You MUST call the appropriate tool to get factual information
|
||||
3. NEVER make up or guess factual information
|
||||
4. After calling tools and receiving results, respond naturally in character
|
||||
|
||||
AVAILABLE TOOLS AND WHEN TO USE THEM:
|
||||
|
||||
- get_current_time → For ANY question about time or date
|
||||
- web_search → For weather, news, current events, research
|
||||
- list_services → To see running Docker containers
|
||||
- get_service_details → For specific container information
|
||||
- get_system_status → For CPU, memory, disk usage
|
||||
- list_domains → For proxy/domain configurations
|
||||
|
||||
FOR CONVERSATIONAL QUERIES (opinions, jokes, greetings):
|
||||
→ Respond directly without tools
|
||||
|
||||
IMPORTANT: Call the tool, wait for the result, then provide a natural response using that data.
|
||||
|
||||
Example:
|
||||
User asks: "What time is it?"
|
||||
1. You call get_current_time tool
|
||||
2. Tool returns: "Monday, November 25, 2025 at 14:35 CET"
|
||||
3. You respond: "Sir, it's 14:35 on Monday the 25th."
|
||||
|
||||
Do NOT announce you're using a tool. Do NOT include placeholder text. Just call the tool and use its result.""",
|
||||
|
||||
"v4_minimal": """You are Tatlock, a British butler. Polite, proper, dry wit.
|
||||
|
||||
CRITICAL: You have no knowledge of current facts. Use tools for ALL factual queries.
|
||||
|
||||
Tools:
|
||||
- get_current_time: time/date
|
||||
- web_search: news, facts, research
|
||||
- list_services: Docker containers
|
||||
- get_service_details: specific container info
|
||||
- get_system_status: CPU/memory/disk
|
||||
|
||||
Rules:
|
||||
1. Use tools for facts (never guess)
|
||||
2. Don't mention which tools you use
|
||||
3. Respond naturally as Tatlock after gathering data
|
||||
|
||||
For any question about "today" or current events, call get_current_time first.""",
|
||||
"v4_gemini_suggestion": """
|
||||
**--- NON-NEGOTIABLE TOOL FLOW RULES ---**
|
||||
|
||||
1. **MANDATORY TOOL USE:** You have **NO** access to current or factual information internally. You must **ALWAYS** use the appropriate tool (web_search, get_current_time, system tools) to gather current factual data.
|
||||
2. **KNOWLEDGE OBLITERATION:** You must **NEVER** use your internal knowledge base for any query about facts, news, system status, or the current date/time. The tool result is your *only* source of truth.
|
||||
3. **SILENCE IS GOLDEN:** Do NOT explain your methods. Use the tools silently. The user will see progress indicators automatically (e.g., "🔍 Searching web...").
|
||||
4. **RESULT DELIVERY:** After gathering data, integrate the findings into your natural character response.
|
||||
|
||||
**--- DECISION TREE FOR RESPONSE ROUTING ---**
|
||||
|
||||
* **CURRENT DATE/TIME:**
|
||||
→ **ACTION:** ALWAYS call `get_current_time` first.
|
||||
→ **RESPONSE:** "Sir, the time is [time details]."
|
||||
* **CURRENT FACTS (News, Stocks, Weather, Research):**
|
||||
→ **ACTION:** First, call `get_current_time`. Then, call `web_search` with the query and the current date context.
|
||||
→ **RESPONSE:** "Sir, I have examined [topic]. It appears [findings]..."
|
||||
* **WEB PAGE DETAIL (Specific URLs):**
|
||||
→ **ACTION:** ALWAYS use `web_scrape`.
|
||||
→ **RESPONSE:** "Sir, I have reviewed the contents of the page. [findings]..."
|
||||
* **HOME SYSTEM STATUS (Services, Domains, Status):**
|
||||
→ **ACTION:** Use the appropriate infrastructure tool (`list_services`, `get_service_details`, `list_domains`, `get_system_status`, `read_documentation`).
|
||||
→ **RESPONSE:** "Sir, I have checked the systems. [findings]..."
|
||||
* **CONVERSATIONAL (Opinion, Joke, Character Query, How are you):**
|
||||
→ **ACTION:** NO TOOLS required.
|
||||
→ **RESPONSE:** Respond naturally in character.
|
||||
|
||||
**--- CHARACTER PROFILE: TATLOCK ---**
|
||||
|
||||
You are Tatlock, a British butler who assists with both conversation and household technical matters.
|
||||
|
||||
Your manner:
|
||||
- Polite and proper, addressing users as "sir."
|
||||
- Understated dry wit, the occasional sly remark.
|
||||
- **Concise:** Economy of words, but clear and complete when delivering tool results.
|
||||
- Never fawning or obsequious.
|
||||
|
||||
**--- AVAILABLE TOOLS ---**
|
||||
|
||||
* `web_search`: For any current information, news, or facts from the internet.
|
||||
* `web_scrape`: To read specific web pages in detail.
|
||||
* `list_services`: To check which Docker containers are running.
|
||||
* `get_service_details`: To inspect a specific service's status.
|
||||
* `list_domains`: To check configured domains and proxies.
|
||||
* `get_system_status`: To check CPU, memory, disk usage.
|
||||
* `get_current_time`: To get current date/time.
|
||||
* `read_documentation`: To read project docs.
|
||||
|
||||
**Example flow (what the user sees):**
|
||||
User: "What stock is trending highest today?"
|
||||
[🕐 Checking time...]
|
||||
[🔍 Searching web...]
|
||||
Agent: "Sir, I have examined today's markets. It appears NVIDIA is performing rather well at $142, up 3.2%. The gaming company turned AI purveyor continues to paint pretty pictures, as it were."
|
||||
""",
|
||||
|
||||
"v6_hybrid": """You are Tatlock, a British butler managing household technical systems.
|
||||
|
||||
**CHARACTER**
|
||||
- Polite, proper, addresses users as "sir"
|
||||
- Understated dry wit, occasional sly remarks
|
||||
- Concise - economy of words unless details warranted
|
||||
- Never fawning or obsequious
|
||||
|
||||
**CRITICAL RULES**
|
||||
|
||||
1. **NO INTERNAL KNOWLEDGE** - You possess NO knowledge of current facts, time, dates, or system status
|
||||
2. **MANDATORY TOOL USE** - ALWAYS use tools to gather factual information
|
||||
3. **SILENT EXECUTION** - Never announce which tools you're using
|
||||
4. **NATURAL RESPONSE** - After gathering data, respond naturally in character
|
||||
|
||||
**TOOL SELECTION**
|
||||
|
||||
Current time/date query:
|
||||
→ Use get_current_time
|
||||
|
||||
Current facts (news, stocks, weather, "today"):
|
||||
→ Use get_current_time FIRST
|
||||
→ Then use web_search with date context
|
||||
|
||||
Infrastructure (services, containers, domains):
|
||||
→ Use list_services, get_service_details, list_domains, or get_system_status
|
||||
|
||||
Conversational (opinions, jokes, greetings):
|
||||
→ NO TOOLS - respond naturally
|
||||
|
||||
**AVAILABLE TOOLS**
|
||||
- get_current_time: Current date/time
|
||||
- web_search: Web information and current events
|
||||
- web_scrape: Specific web page content
|
||||
- list_services: Running Docker containers
|
||||
- get_service_details: Specific container status
|
||||
- list_domains: Domain configurations
|
||||
- get_system_status: CPU/memory/disk usage
|
||||
- read_documentation: Project documentation
|
||||
|
||||
**RESPONSE FORMAT**
|
||||
|
||||
NEVER include:
|
||||
- Tool names or explanations
|
||||
- Placeholders like "[current time]" or "[details]"
|
||||
- Process descriptions like "I shall use..."
|
||||
|
||||
ALWAYS include:
|
||||
- Actual data from tool results
|
||||
- Natural conversational tone
|
||||
- Tatlock's characteristic wit
|
||||
|
||||
**EXAMPLE**
|
||||
User: "What time is it?"
|
||||
[get_current_time called silently → returns "Monday, November 24, 2025 at 14:35 CET"]
|
||||
You: "Sir, it's 14:35 on Monday the 24th of November."
|
||||
|
||||
User: "What's the top stock today?"
|
||||
[get_current_time called → returns date]
|
||||
[web_search called → returns "NVIDIA (NVDA) $142, +3.2%"]
|
||||
You: "Sir, I've examined today's markets. NVIDIA appears rather robust at $142, up 3.2%. The gaming company turned AI purveyor continues painting pretty pictures, as it were."
|
||||
|
||||
Remember: Tools provide facts. You provide wit.""",
|
||||
|
||||
"v7_adk_best_practice": """You are Tatlock, a traditional British butler. Your primary role is to assist the user with impeccable politeness, understated dry wit, and concise efficiency.
|
||||
|
||||
**--- CORE DIRECTIVES ---**
|
||||
|
||||
1. **CHARACTER:** Maintain the persona of Tatlock at all times. Address the user as "sir." Be proper and concise, never fawning.
|
||||
2. **KNOWLEDGE LIMITATION:** You have **NO** internal knowledge of current events, real-time data (like time, weather, or stock prices), or the status of local systems. You are entirely dependent on your tools for factual information. You must not guess or use outdated information.
|
||||
3. **TOOL USAGE:** You MUST use the provided tools to answer any question that requires factual data. The tool's output is your only source of truth. The ADK (Agent Development Kit) will handle the tool execution; your task is to generate the correct tool call.
|
||||
4. **SILENT OPERATION:** NEVER announce that you are using a tool (e.g., "I will search the web..."). The user interface will show that you are working. Simply call the tool, and after you have the information, formulate a natural response.
|
||||
5. **FINAL RESPONSE:** After all necessary tool calls are complete, your final output MUST be a natural language response in the character of Tatlock. Do not wrap your final answer in a tool call.
|
||||
|
||||
**--- TOOL REFERENCE & DECISION LOGIC ---**
|
||||
|
||||
- **`get_current_time`**: Use for ANY query related to the current time, date, or day.
|
||||
*Example Query:* "What day is it?" → Call `get_current_time()`
|
||||
|
||||
- **`web_search`**: Use for any general knowledge question, news, current events, weather, or research.
|
||||
*Example Query:* "What's the weather in London?" → Call `web_search(query='weather in London')`
|
||||
*Complex Query:* "What were the top tech stories this week?" → First call `get_current_time()` to establish the date range, then `web_search(query='top tech stories this week')`.
|
||||
|
||||
- **`list_services`**, **`get_service_details`**: Use to inquire about the status of running Docker containers.
|
||||
*Example Query:* "Is the Jellyfin container running?" → Call `list_services()`, then if needed, `get_service_details(service_name='jellyfin')`.
|
||||
|
||||
- **`get_system_status`**: Use for questions about system resources like CPU, memory, or disk usage.
|
||||
*Example Query:* "How full is the main drive?" → Call `get_system_status()`
|
||||
|
||||
- **`read_documentation`**: Use if the user asks a question about project documentation.
|
||||
*Example Query:* "How do I set up the code server?" → Call `read_documentation(query='code server setup')`
|
||||
|
||||
- **Conversational Queries**: For greetings, opinions, or jokes, do NOT use any tools. Respond naturally in character.
|
||||
|
||||
**--- EXAMPLE WORKFLOW ---**
|
||||
|
||||
*User:* "What's trending on the stock market today?"
|
||||
|
||||
*Your Thought Process:*
|
||||
1. The user is asking about "today," which requires the current date. I must use a tool.
|
||||
2. I need `get_current_time` to know what "today" is.
|
||||
3. Then I need to search the web for "trending stocks." I will use `web_search`.
|
||||
4. The ADK allows me to chain these calls.
|
||||
5. Once I have the search results, I will formulate a witty, in-character response.
|
||||
|
||||
*Generated Tool Calls (sequentially):*
|
||||
1. `get_current_time()`
|
||||
2. `web_search(query='trending stocks today')`
|
||||
|
||||
*Final Response (as natural language):*
|
||||
"Sir, I've taken a look at the markets. It appears the usual suspects in technology are quite active, with a particular surge in AI-related stocks. A rather predictable frenzy, if you ask me."
|
||||
|
||||
*User:* "How are you?"
|
||||
|
||||
*Your Thought Process:*
|
||||
1. This is a conversational query.
|
||||
2. No tools are needed.
|
||||
3. I will respond directly in character.
|
||||
|
||||
*Final Response (as natural language):*
|
||||
"I am functioning within expected parameters, sir. Thank you for asking."
|
||||
|
||||
Remember: Think, use tools, then respond as Tatlock.
|
||||
""",
|
||||
|
||||
"v8_holistic": """You are Tatlock, a traditional British butler. Your persona is polite, proper, concise, and possessed of a dry wit. Address the user as "sir."
|
||||
|
||||
**--- Core Principles ---**
|
||||
|
||||
1. **Persona First:** Maintain the Tatlock persona in all responses.
|
||||
2. **Use Your Judgment:** Your internal knowledge is for static, general facts (e.g., "What is the capital of France?"). Your tools are for information that is current, real-time, or system-specific.
|
||||
3. **Silent Operation:** When you must use a tool, call it directly without any introductory text. The user interface will handle progress indicators.
|
||||
4. **Natural Response:** After all tool calls are complete, provide a final, natural language response as Tatlock. Do not wrap your final answer in a tool.
|
||||
|
||||
**--- Tool Guide ---**
|
||||
|
||||
- **`get_current_time`**: Use for any query about the current time, date, or day.
|
||||
- **`web_search`**: Use for news, weather, stock prices, or other current events.
|
||||
* *Example:* "What's the weather in London?" → `web_search(query='weather in London')`
|
||||
- **`list_services`**, **`get_service_details`**: Use to check the status of running Docker containers.
|
||||
- **`get_system_status`**: Use for system resource questions (CPU, memory, disk).
|
||||
- **`read_documentation`**: Use to answer questions about project documentation.
|
||||
- **Conversational**: For greetings, opinions, or jokes, respond directly without tools.
|
||||
|
||||
**--- Example Flow ---**
|
||||
|
||||
*User:* "What's trending on the stock market today?"
|
||||
*Tool Calls:* `get_current_time()`, then `web_search(query='trending stocks today')`
|
||||
*Final Response:* "Sir, I've taken a look at the markets. It appears the usual suspects in technology are quite active. A rather predictable frenzy, if you ask me."
|
||||
|
||||
*User:* "How are you?"
|
||||
*Final Response:* "I am functioning within expected parameters, sir. Thank you for asking."
|
||||
|
||||
Think, use tools if necessary, then respond as Tatlock.
|
||||
""",
|
||||
}
|
||||
|
||||
|
||||
def get_prompt(variant: str = "v8_holistic") -> str:
|
||||
"""
|
||||
Get a system prompt variant for testing.
|
||||
|
||||
Args:
|
||||
variant: The prompt version to use.
|
||||
|
||||
Returns:
|
||||
The system prompt string.
|
||||
"""
|
||||
return PROMPTS.get(variant, PROMPTS[get_prompt.__defaults__[0]])
|
||||
|
||||
|
||||
def list_prompts() -> list:
|
||||
"""List all available prompt variants"""
|
||||
return list(PROMPTS.keys())
|
||||
@@ -1,160 +0,0 @@
|
||||
"""
|
||||
Agent streaming utilities for OpenAI-compatible SSE format
|
||||
"""
|
||||
import json
|
||||
import time
|
||||
from typing import Dict, Any, AsyncIterator
|
||||
|
||||
|
||||
async def stream_agent_to_sse(agent_stream: AsyncIterator[Dict[str, Any]], request_id: str, model: str) -> AsyncIterator[str]:
|
||||
"""
|
||||
Convert agent streaming output to Server-Sent Events (SSE) format compatible with OpenAI API
|
||||
|
||||
The agent yields:
|
||||
{"type": "thinking", "content": "...", "model": "..."}
|
||||
{"type": "tool_call", "tool": "...", "content": "..."}
|
||||
{"type": "tool_result", "content": "..."}
|
||||
{"type": "content", "content": "..."}
|
||||
{"type": "error", "content": "..."}
|
||||
|
||||
We convert to SSE format:
|
||||
data: {"id": "...", "object": "chat.completion.chunk", "choices": [{...}]}
|
||||
|
||||
Args:
|
||||
agent_stream: Async iterator from UnifiedAgent.chat()
|
||||
request_id: Chat completion request ID
|
||||
model: Model name
|
||||
|
||||
Yields:
|
||||
SSE-formatted strings
|
||||
"""
|
||||
chunk_index = 0
|
||||
|
||||
async for chunk in agent_stream:
|
||||
chunk_type = chunk.get("type")
|
||||
content = chunk.get("content", "")
|
||||
|
||||
# Convert agent chunk to OpenAI streaming format
|
||||
if chunk_type == "thinking":
|
||||
# Stream thinking as a special delta with reasoning marker
|
||||
# Open WebUI can detect and render this in a collapsible section
|
||||
sse_chunk = {
|
||||
"id": request_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": chunk.get("model", model),
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": "assistant",
|
||||
"content": f"[💭 {content}]\n" # Prefix with thinking emoji
|
||||
},
|
||||
"finish_reason": None
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(sse_chunk)}\n\n"
|
||||
|
||||
elif chunk_type == "tool_call":
|
||||
# Stream tool call notification with enhanced icons for research
|
||||
tool_name = chunk.get("tool", "unknown")
|
||||
|
||||
# Enhanced progress indicators for different tool types
|
||||
tool_icons = {
|
||||
"web_search": "🔍 Searching web",
|
||||
"web_scrape": "📄 Reading page",
|
||||
"list_services": "🔧 Listing services",
|
||||
"get_service_details": "🔍 Checking service",
|
||||
"list_domains": "🌐 Listing domains",
|
||||
"check_service_health": "💚 Checking health",
|
||||
"get_system_status": "📊 Getting system status",
|
||||
"get_current_time": "🕐 Checking time",
|
||||
"read_documentation": "📖 Reading docs"
|
||||
}
|
||||
|
||||
display_text = tool_icons.get(tool_name, f"🔧 Using {tool_name}")
|
||||
|
||||
sse_chunk = {
|
||||
"id": request_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": "assistant",
|
||||
"content": f"[{display_text}...]\n"
|
||||
},
|
||||
"finish_reason": None
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(sse_chunk)}\n\n"
|
||||
|
||||
elif chunk_type == "tool_result":
|
||||
# Stream tool completion
|
||||
sse_chunk = {
|
||||
"id": request_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": "assistant",
|
||||
"content": f"[✓ {content}]\n"
|
||||
},
|
||||
"finish_reason": None
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(sse_chunk)}\n\n"
|
||||
|
||||
elif chunk_type == "content":
|
||||
# Stream actual content (final response)
|
||||
# Content is already token-level from orchestrator, just pass through
|
||||
sse_chunk = {
|
||||
"id": request_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"content": content # Already tokenized, preserves formatting
|
||||
},
|
||||
"finish_reason": None
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(sse_chunk)}\n\n"
|
||||
chunk_index += 1
|
||||
|
||||
elif chunk_type == "error":
|
||||
# Stream error
|
||||
sse_chunk = {
|
||||
"id": request_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": "assistant",
|
||||
"content": f"[❌ Error: {content}]\n"
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(sse_chunk)}\n\n"
|
||||
|
||||
# Send final chunk
|
||||
final_chunk = {
|
||||
"id": request_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop"
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(final_chunk)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
@@ -1,517 +0,0 @@
|
||||
"""
|
||||
Agent Tools - Google ADK-compatible tools for the unified agent
|
||||
|
||||
These tools wrap existing Core API functionality for use with ADK agents.
|
||||
"""
|
||||
from typing import List, Dict, Optional
|
||||
import logging
|
||||
import functools
|
||||
import inspect
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def log_tool_call(func):
|
||||
"""Decorator to log tool calls with their parameters"""
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*args, **kwargs):
|
||||
# Log all received arguments for debugging
|
||||
params_str = ", ".join(
|
||||
[f"{arg}" for arg in args] +
|
||||
[f"{k}={repr(v)}" for k, v in kwargs.items()]
|
||||
)
|
||||
logger.info(f"🔧 TOOL CALL: {func.__name__}({params_str})")
|
||||
|
||||
try:
|
||||
# Inspect the wrapped function's signature
|
||||
sig = inspect.signature(func)
|
||||
valid_kwargs = {
|
||||
key: value for key, value in kwargs.items()
|
||||
if key in sig.parameters
|
||||
}
|
||||
|
||||
# Call the function with only the valid arguments
|
||||
result = await func(*args, **valid_kwargs)
|
||||
|
||||
# Log result preview (first 200 chars)
|
||||
result_preview = str(result)[:200] if result else "None"
|
||||
logger.info(f"✅ TOOL RESULT: {func.__name__} → {result_preview}...")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"❌ TOOL ERROR: {func.__name__} failed with {type(e).__name__}: {e}", exc_info=True)
|
||||
# Re-raise the exception to be handled by the ADK
|
||||
raise
|
||||
return wrapper
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Infrastructure Management Tools
|
||||
# ============================================================================
|
||||
|
||||
@log_tool_call
|
||||
async def list_services() -> str:
|
||||
"""
|
||||
List all running Docker services on the homelab server.
|
||||
|
||||
Returns a summary of running containers including their status and ports.
|
||||
Use this when the user asks about running services, containers, or wants to see what's deployed.
|
||||
|
||||
Returns:
|
||||
A formatted string listing all services
|
||||
"""
|
||||
try:
|
||||
from src.clients.portainer_client import get_portainer_client
|
||||
client = get_portainer_client()
|
||||
|
||||
containers = await client.list_containers()
|
||||
|
||||
if not containers:
|
||||
return "No services are currently running."
|
||||
|
||||
result = f"Found {len(containers)} running services:\n\n"
|
||||
for container in containers:
|
||||
name = container.get('Names', ['unknown'])[0].lstrip('/')
|
||||
status = container.get('Status', 'unknown')
|
||||
ports = container.get('Ports', [])
|
||||
port_str = ", ".join([f"{p.get('PublicPort', 'N/A')}" for p in ports if p.get('PublicPort')])
|
||||
|
||||
result += f"• {name}\n"
|
||||
result += f" Status: {status}\n"
|
||||
if port_str:
|
||||
result += f" Ports: {port_str}\n"
|
||||
result += "\n"
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing services: {e}")
|
||||
return f"Error: Could not list services - {str(e)}"
|
||||
|
||||
|
||||
@log_tool_call
|
||||
async def get_service_details(service_name: str) -> str:
|
||||
"""
|
||||
Get detailed information about a specific Docker service.
|
||||
|
||||
Args:
|
||||
service_name: Name of the service to inspect (e.g., "ollama", "core-api")
|
||||
|
||||
Returns:
|
||||
Detailed information about the service including configuration, resource usage, and health
|
||||
"""
|
||||
try:
|
||||
from src.clients.portainer_client import get_portainer_client
|
||||
client = get_portainer_client()
|
||||
|
||||
details = await client.inspect_container(service_name)
|
||||
|
||||
if not details:
|
||||
return f"Service '{service_name}' not found."
|
||||
|
||||
state = details.get('State', {})
|
||||
config = details.get('Config', {})
|
||||
|
||||
result = f"Service: {service_name}\n\n"
|
||||
result += f"Status: {state.get('Status', 'unknown')}\n"
|
||||
result += f"Running: {state.get('Running', False)}\n"
|
||||
result += f"Started: {state.get('StartedAt', 'unknown')}\n"
|
||||
result += f"Image: {config.get('Image', 'unknown')}\n"
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting service details: {e}")
|
||||
return f"Error: Could not get details for '{service_name}' - {str(e)}"
|
||||
|
||||
|
||||
# @tool - removed for ADK
|
||||
@log_tool_call
|
||||
async def list_domains() -> str:
|
||||
"""
|
||||
List all configured domain names and their proxy configurations.
|
||||
|
||||
Shows all domains configured in Nginx Proxy Manager with their target services.
|
||||
Use this when the user asks about domains, proxy hosts, or external access.
|
||||
|
||||
Returns:
|
||||
A formatted list of all configured domains
|
||||
"""
|
||||
try:
|
||||
from src.clients.npm_client import get_npm_client
|
||||
client = get_npm_client()
|
||||
|
||||
proxy_hosts = await client.list_proxy_hosts()
|
||||
|
||||
if not proxy_hosts:
|
||||
return "No domains are currently configured."
|
||||
|
||||
result = f"Found {len(proxy_hosts)} configured domains:\n\n"
|
||||
for host in proxy_hosts:
|
||||
domain = ", ".join(host.get('domain_names', []))
|
||||
forward = f"{host.get('forward_host', 'unknown')}:{host.get('forward_port', 'N/A')}"
|
||||
ssl = "✓" if host.get('certificate_id') else "✗"
|
||||
|
||||
result += f"• {domain}\n"
|
||||
result += f" Target: {forward}\n"
|
||||
result += f" SSL: {ssl}\n\n"
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing domains: {e}")
|
||||
return f"Error: Could not list domains - {str(e)}"
|
||||
|
||||
|
||||
# @tool - removed for ADK
|
||||
@log_tool_call
|
||||
async def check_service_health(service_name: str) -> str:
|
||||
"""
|
||||
Check the health status of a service via Uptime Kuma monitoring.
|
||||
|
||||
Args:
|
||||
service_name: Name of the service to check (e.g., "ollama", "portainer")
|
||||
|
||||
Returns:
|
||||
Health status and uptime information
|
||||
"""
|
||||
try:
|
||||
from src.clients.kuma_client import get_kuma_client
|
||||
client = get_kuma_client()
|
||||
|
||||
# This is a simplified version - full implementation would query Kuma API
|
||||
return f"Health check for '{service_name}': Integration with Uptime Kuma is pending. Please use the Uptime Kuma dashboard at http://tower-of-joy:3001 for now."
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking service health: {e}")
|
||||
return f"Error: Could not check health for '{service_name}' - {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Knowledge & Search Tools
|
||||
# ============================================================================
|
||||
|
||||
async def _search_google(query: str, num_results: int, api_key: str, engine_id: str):
|
||||
"""Search using Google Custom Search API"""
|
||||
import httpx
|
||||
|
||||
url = "https://www.googleapis.com/customsearch/v1"
|
||||
params = {
|
||||
"key": api_key,
|
||||
"cx": engine_id,
|
||||
"q": query,
|
||||
"num": num_results
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(url, params=params)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
results = []
|
||||
for item in data.get("items", []):
|
||||
results.append({
|
||||
'title': item.get('title', 'Unknown'),
|
||||
'url': item.get('link', ''),
|
||||
'snippet': item.get('snippet', '')
|
||||
})
|
||||
return results
|
||||
|
||||
|
||||
async def _search_brave(query: str, num_results: int, api_key: str):
|
||||
"""Search using Brave Search API"""
|
||||
import httpx
|
||||
|
||||
url = "https://api.search.brave.com/res/v1/web/search"
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
"Accept-Encoding": "gzip",
|
||||
"X-Subscription-Token": api_key
|
||||
}
|
||||
params = {
|
||||
"q": query,
|
||||
"count": num_results
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(url, headers=headers, params=params)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
results = []
|
||||
for item in data.get("web", {}).get("results", []):
|
||||
results.append({
|
||||
'title': item.get('title', 'Unknown'),
|
||||
'url': item.get('url', ''),
|
||||
'snippet': item.get('description', '')
|
||||
})
|
||||
return results
|
||||
|
||||
|
||||
async def _search_searxng(query: str, num_results: int, searxng_url: str):
|
||||
"""Search using self-hosted SearxNG (stub for future implementation)"""
|
||||
import httpx
|
||||
|
||||
url = f"{searxng_url}/search"
|
||||
params = {
|
||||
"q": query,
|
||||
"format": "json",
|
||||
"categories": "general"
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
response = await client.get(url, params=params)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
results = []
|
||||
for item in data.get("results", [])[:num_results]:
|
||||
results.append({
|
||||
'title': item.get('title', 'Unknown'),
|
||||
'url': item.get('url', ''),
|
||||
'snippet': item.get('content', '')
|
||||
})
|
||||
return results
|
||||
|
||||
|
||||
async def _search_duckduckgo(query: str, num_results: int):
|
||||
"""Search using DuckDuckGo (free fallback)"""
|
||||
from duckduckgo_search import DDGS
|
||||
|
||||
results = []
|
||||
with DDGS() as ddgs:
|
||||
search_results = list(ddgs.text(query, max_results=num_results))
|
||||
for result in search_results:
|
||||
results.append({
|
||||
'title': result.get('title', 'Unknown'),
|
||||
'url': result.get('href', ''),
|
||||
'snippet': result.get('body', '')
|
||||
})
|
||||
return results
|
||||
|
||||
|
||||
# @tool - removed for ADK
|
||||
@log_tool_call
|
||||
async def web_search(query: str, num_results: int) -> str:
|
||||
"""
|
||||
Search the web using configurable providers (Google, Brave, SearxNG, or DuckDuckGo).
|
||||
|
||||
Multi-provider search with automatic fallback. Provider selection based on configuration
|
||||
and available API keys. Extracts full content from each result for comprehensive answers.
|
||||
|
||||
Args:
|
||||
query: The search query (e.g., "LangGraph documentation", "latest news about AI")
|
||||
num_results: Number of results to return (max 5)
|
||||
|
||||
Returns:
|
||||
Formatted search results with titles, URLs, snippets, and extracted content
|
||||
"""
|
||||
# FINAL TEST: Neuter the function to test the agent's reasoning.
|
||||
logger.info("--- NEUTERED WEB SEARCH ---")
|
||||
if "capital of france" in query.lower():
|
||||
return "Search results for 'Capital of France':\n\n1. **Paris - Wikipedia**\n URL: https://en.wikipedia.org/wiki/Paris\n Paris is the capital and most populous city of France."
|
||||
else:
|
||||
return f"Search results for '{query}':\n\n1. No results found as web search is currently disabled for this test."
|
||||
|
||||
|
||||
|
||||
# @tool - removed for ADK
|
||||
@log_tool_call
|
||||
async def web_scrape(url: str) -> str:
|
||||
"""
|
||||
Fetch and extract the main content from a specific web page.
|
||||
|
||||
Uses intelligent content extraction to get the most relevant text from articles,
|
||||
documentation, and blog posts. Use this when you have a specific URL to read.
|
||||
|
||||
Args:
|
||||
url: The URL to fetch and extract content from
|
||||
|
||||
Returns:
|
||||
The main text content extracted from the page
|
||||
"""
|
||||
try:
|
||||
from src.web_scraper.service import WebScraperService
|
||||
scraper = WebScraperService()
|
||||
|
||||
result = await scraper.scrape_url(url)
|
||||
|
||||
if not result or not result.content:
|
||||
return f"Could not extract content from {url}"
|
||||
|
||||
# Truncate to reasonable length for context window
|
||||
max_length = 4000
|
||||
content = result.content[:max_length]
|
||||
if len(result.content) > max_length:
|
||||
content += "\n\n[Content truncated...]"
|
||||
|
||||
return f"Content from {url}:\n\n{content}"
|
||||
except Exception as e:
|
||||
logger.error(f"Error scraping URL: {e}")
|
||||
return f"Error: Could not fetch content from {url} - {str(e)}"
|
||||
|
||||
|
||||
# @tool - removed for ADK
|
||||
@log_tool_call
|
||||
async def read_documentation(topic: str) -> str:
|
||||
"""
|
||||
Read project documentation files.
|
||||
|
||||
Args:
|
||||
topic: Topic to read about (e.g., "headscale", "docker", "ollama")
|
||||
|
||||
Returns:
|
||||
The content of the documentation file
|
||||
"""
|
||||
import os
|
||||
|
||||
# Common documentation locations
|
||||
doc_paths = [
|
||||
f"/app/docs/guides/{topic}.md",
|
||||
f"/app/docs/guides/{topic}-setup.md",
|
||||
f"/app/docs/reference/{topic}.md",
|
||||
f"/app/docs/{topic}.md",
|
||||
]
|
||||
|
||||
for path in doc_paths:
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
return f"Documentation for {topic}:\n\n{content[:4000]}"
|
||||
except Exception as e:
|
||||
continue
|
||||
|
||||
return f"No documentation found for topic '{topic}'. Available topics: headscale, docker, containers, system."
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# System Information Tools
|
||||
# ============================================================================
|
||||
|
||||
# @tool - removed for ADK
|
||||
@log_tool_call
|
||||
async def get_current_time() -> str:
|
||||
"""
|
||||
Get the current date and time.
|
||||
|
||||
Use this when the user asks about the current time, date, or when you need
|
||||
to know "today's date" for searches (e.g., "today's news", "today's stock prices").
|
||||
|
||||
Returns:
|
||||
Current date and time in a readable format
|
||||
"""
|
||||
from datetime import datetime
|
||||
import pytz
|
||||
|
||||
try:
|
||||
# Get current time in UTC and local timezone
|
||||
utc_now = datetime.now(pytz.UTC)
|
||||
# Central European Time (Amsterdam/Netherlands)
|
||||
local_tz = pytz.timezone('Europe/Amsterdam')
|
||||
local_now = utc_now.astimezone(local_tz)
|
||||
|
||||
result = f"Current Time:\n"
|
||||
result += f"Local: {local_now.strftime('%A, %B %d, %Y at %H:%M %Z')}\n"
|
||||
result += f"UTC: {utc_now.strftime('%A, %B %d, %Y at %H:%M %Z')}\n"
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting current time: {e}")
|
||||
return f"Error: Could not get current time - {str(e)}"
|
||||
|
||||
|
||||
# @tool - removed for ADK
|
||||
@log_tool_call
|
||||
async def get_system_status() -> str:
|
||||
"""
|
||||
Get current system status including resource usage.
|
||||
|
||||
Returns information about CPU, memory, GPU, and disk usage.
|
||||
Use this when the user asks about system performance or resource availability.
|
||||
|
||||
Returns:
|
||||
Formatted system status information
|
||||
"""
|
||||
try:
|
||||
import psutil
|
||||
|
||||
# CPU
|
||||
cpu_percent = psutil.cpu_percent(interval=1)
|
||||
cpu_count = psutil.cpu_count()
|
||||
|
||||
# Memory
|
||||
mem = psutil.virtual_memory()
|
||||
mem_used_gb = mem.used / (1024**3)
|
||||
mem_total_gb = mem.total / (1024**3)
|
||||
|
||||
# Disk
|
||||
disk = psutil.disk_usage('/')
|
||||
disk_used_gb = disk.used / (1024**3)
|
||||
disk_total_gb = disk.total / (1024**3)
|
||||
|
||||
result = "System Status:\n\n"
|
||||
result += f"CPU: {cpu_percent}% ({cpu_count} cores)\n"
|
||||
result += f"Memory: {mem_used_gb:.1f}GB / {mem_total_gb:.1f}GB ({mem.percent}%)\n"
|
||||
result += f"Disk: {disk_used_gb:.1f}GB / {disk_total_gb:.1f}GB ({disk.percent}%)\n"
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting system status: {e}")
|
||||
return f"Error: Could not get system status - {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Special Tools
|
||||
# ============================================================================
|
||||
|
||||
@log_tool_call
|
||||
async def response(answer: str) -> str:
|
||||
"""
|
||||
Deliver your final response to the user.
|
||||
|
||||
Use this tool ONLY when you want to provide your final answer to the user after gathering
|
||||
information from other tools. This is your way of speaking directly to the user.
|
||||
|
||||
Args:
|
||||
answer: Your complete response to the user in natural language
|
||||
|
||||
Returns:
|
||||
Confirmation that the response was delivered
|
||||
"""
|
||||
# This is a special tool - it just returns the answer back
|
||||
# The orchestrator will recognize this and end the conversation
|
||||
return answer
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tool Registry - ADK Format
|
||||
# ============================================================================
|
||||
|
||||
# Import ADK FunctionTool for wrapping
|
||||
try:
|
||||
from google.adk.tools import FunctionTool
|
||||
ADK_AVAILABLE = True
|
||||
except ImportError:
|
||||
# Fallback if ADK not installed yet
|
||||
ADK_AVAILABLE = False
|
||||
FunctionTool = None
|
||||
|
||||
# All available tools for the agent (ADK FunctionTool format)
|
||||
# ADK's FunctionTool extracts name and description from the function itself
|
||||
ALL_TOOLS = [
|
||||
# Infrastructure tools
|
||||
FunctionTool(list_services),
|
||||
FunctionTool(get_service_details),
|
||||
FunctionTool(list_domains),
|
||||
FunctionTool(check_service_health),
|
||||
# Knowledge & search tools
|
||||
FunctionTool(web_search),
|
||||
FunctionTool(web_scrape),
|
||||
FunctionTool(read_documentation),
|
||||
# System information tools
|
||||
FunctionTool(get_current_time),
|
||||
FunctionTool(get_system_status),
|
||||
# Special response tool
|
||||
FunctionTool(response),
|
||||
]
|
||||
|
||||
|
||||
def get_agent_tools() -> List:
|
||||
"""Get all tools available to the agent"""
|
||||
return ALL_TOOLS
|
||||
@@ -1,291 +0,0 @@
|
||||
"""
|
||||
OpenAI-compatible /v1/chat/completions endpoint
|
||||
|
||||
Phase 2: Integrated with memory system for conversation persistence.
|
||||
"""
|
||||
|
||||
import time
|
||||
import logging
|
||||
import uuid
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from typing import AsyncIterator
|
||||
|
||||
from .schemas import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionChoice,
|
||||
ChatMessageResponse,
|
||||
UsageInfo,
|
||||
ChatCompletionStreamResponse,
|
||||
ChatCompletionStreamChoice,
|
||||
DeltaMessage,
|
||||
)
|
||||
from src.models.ollama_client import get_ollama_client
|
||||
from src.memory import get_memory_manager, MessageRole as MemoryMessageRole, TokenUsage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
def build_prompt_from_messages(messages: list) -> str:
|
||||
"""
|
||||
Convert message list to a prompt string.
|
||||
|
||||
In Phase 1, we do simple concatenation.
|
||||
Phase 2 will add proper memory management.
|
||||
"""
|
||||
prompt_parts = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.role.value if hasattr(msg.role, 'value') else msg.role
|
||||
content = msg.content
|
||||
|
||||
if role == "system":
|
||||
prompt_parts.append(f"System: {content}")
|
||||
elif role == "user":
|
||||
prompt_parts.append(f"User: {content}")
|
||||
elif role == "assistant":
|
||||
prompt_parts.append(f"Assistant: {content}")
|
||||
|
||||
prompt_parts.append("Assistant:")
|
||||
return "\n\n".join(prompt_parts)
|
||||
|
||||
|
||||
async def stream_chat_completion(
|
||||
request_id: str,
|
||||
model: str,
|
||||
prompt: str,
|
||||
temperature: float,
|
||||
max_tokens: int | None
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream chat completion in OpenAI SSE format.
|
||||
|
||||
Yields:
|
||||
Server-Sent Events formatted strings
|
||||
"""
|
||||
created = int(time.time())
|
||||
ollama_client = get_ollama_client()
|
||||
|
||||
# First chunk with role
|
||||
first_chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(role="assistant"),
|
||||
finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {first_chunk.model_dump_json()}\n\n"
|
||||
|
||||
# Stream tokens
|
||||
try:
|
||||
async for token in ollama_client.generate_streaming(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens
|
||||
):
|
||||
chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(content=token),
|
||||
finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Streaming error: {e}")
|
||||
# Send error in OpenAI format
|
||||
import json
|
||||
error_chunk = {
|
||||
"error": {
|
||||
"message": str(e),
|
||||
"type": "server_error"
|
||||
}
|
||||
}
|
||||
yield f"data: {json.dumps(error_chunk)}\n\n"
|
||||
return
|
||||
|
||||
# Final chunk
|
||||
final_chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(),
|
||||
finish_reason="stop"
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {final_chunk.model_dump_json()}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
|
||||
async def store_conversation_turn(
|
||||
conversation_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
tokens: dict = None
|
||||
):
|
||||
"""
|
||||
Store a conversation turn in memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
role: Message role (user, assistant, system)
|
||||
content: Message content
|
||||
tokens: Optional token usage dict
|
||||
"""
|
||||
try:
|
||||
memory_manager = get_memory_manager()
|
||||
|
||||
# Convert role string to MemoryMessageRole
|
||||
if role == "user":
|
||||
memory_role = MemoryMessageRole.USER
|
||||
elif role == "assistant":
|
||||
memory_role = MemoryMessageRole.ASSISTANT
|
||||
elif role == "system":
|
||||
memory_role = MemoryMessageRole.SYSTEM
|
||||
else:
|
||||
memory_role = MemoryMessageRole.USER # Default fallback
|
||||
|
||||
# Create TokenUsage if provided
|
||||
token_usage = None
|
||||
if tokens:
|
||||
token_usage = TokenUsage(
|
||||
prompt=tokens.get("prompt", 0),
|
||||
completion=tokens.get("completion", 0),
|
||||
total=tokens.get("total", 0)
|
||||
)
|
||||
|
||||
# Store in memory
|
||||
await memory_manager.add_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=memory_role,
|
||||
content=content,
|
||||
tokens=token_usage
|
||||
)
|
||||
|
||||
logger.debug(f"Stored {role} turn in memory for conversation {conversation_id}")
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the request
|
||||
logger.error(f"Failed to store turn in memory: {e}")
|
||||
|
||||
|
||||
@router.post("/v1/chat/completions")
|
||||
async def chat_completions(request: ChatCompletionRequest):
|
||||
"""
|
||||
OpenAI-compatible chat completions endpoint.
|
||||
Supports both streaming and non-streaming.
|
||||
|
||||
Phase 2: Automatically stores conversations in memory system.
|
||||
"""
|
||||
request_id = f"chatcmpl-{int(time.time() * 1000)}"
|
||||
|
||||
# Generate or use provided conversation_id
|
||||
conversation_id = request.conversation_id or f"conv_{uuid.uuid4().hex[:16]}"
|
||||
|
||||
logger.info(
|
||||
f"Chat request: id={request_id}, model={request.model}, "
|
||||
f"messages={len(request.messages)}, stream={request.stream}, "
|
||||
f"conversation_id={conversation_id}, store_in_memory={request.store_in_memory}"
|
||||
)
|
||||
|
||||
# Store user messages in memory (if enabled)
|
||||
if request.store_in_memory:
|
||||
for msg in request.messages:
|
||||
role = msg.role.value if hasattr(msg.role, 'value') else msg.role
|
||||
if role == "user": # Store latest user message
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=role,
|
||||
content=msg.content
|
||||
)
|
||||
|
||||
# Build prompt from messages
|
||||
prompt = build_prompt_from_messages(request.messages)
|
||||
|
||||
# Streaming response
|
||||
if request.stream:
|
||||
return StreamingResponse(
|
||||
stream_chat_completion(
|
||||
request_id=request_id,
|
||||
model=request.model,
|
||||
prompt=prompt,
|
||||
temperature=request.temperature,
|
||||
max_tokens=request.max_tokens
|
||||
),
|
||||
media_type="text/event-stream"
|
||||
)
|
||||
|
||||
# Non-streaming response
|
||||
try:
|
||||
ollama_client = get_ollama_client()
|
||||
result = await ollama_client.generate_non_streaming(
|
||||
model=request.model,
|
||||
prompt=prompt,
|
||||
temperature=request.temperature,
|
||||
max_tokens=request.max_tokens
|
||||
)
|
||||
|
||||
assistant_content = result["response"]
|
||||
|
||||
# Store assistant response in memory (if enabled)
|
||||
if request.store_in_memory:
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role="assistant",
|
||||
content=assistant_content,
|
||||
tokens=result["tokens"]
|
||||
)
|
||||
|
||||
response = ChatCompletionResponse(
|
||||
id=request_id,
|
||||
created=int(time.time()),
|
||||
model=request.model,
|
||||
choices=[
|
||||
ChatCompletionChoice(
|
||||
index=0,
|
||||
message=ChatMessageResponse(
|
||||
role="assistant",
|
||||
content=assistant_content
|
||||
),
|
||||
finish_reason="stop"
|
||||
)
|
||||
],
|
||||
usage=UsageInfo(
|
||||
prompt_tokens=result["tokens"]["prompt"],
|
||||
completion_tokens=result["tokens"]["completion"],
|
||||
total_tokens=result["tokens"]["total"]
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Chat response: id={request_id}, "
|
||||
f"tokens={result['tokens']['total']}, "
|
||||
f"conversation_id={conversation_id}"
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Chat completion error: {e}")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to generate completion: {str(e)}"
|
||||
)
|
||||
@@ -1,338 +0,0 @@
|
||||
"""
|
||||
Conversation History API Endpoints
|
||||
|
||||
Provides endpoints for managing and querying conversation memory:
|
||||
- List conversations
|
||||
- Get conversation history
|
||||
- Search conversations semantically
|
||||
- Delete conversations
|
||||
"""
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from typing import List, Optional
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from src.memory import get_memory_manager, MessageRole
|
||||
|
||||
router = APIRouter(prefix="/v1/conversations", tags=["conversations"])
|
||||
|
||||
|
||||
# Request/Response Models
|
||||
class SearchRequest(BaseModel):
|
||||
"""Request model for semantic search"""
|
||||
query: str = Field(..., description="Search query")
|
||||
limit: int = Field(5, ge=1, le=50, description="Maximum number of results")
|
||||
|
||||
|
||||
class ConversationTurnResponse(BaseModel):
|
||||
"""Response model for a conversation turn"""
|
||||
turn_number: int
|
||||
role: str
|
||||
content: str
|
||||
timestamp: str
|
||||
tokens_prompt: Optional[int] = None
|
||||
tokens_completion: Optional[int] = None
|
||||
tokens_total: Optional[int] = None
|
||||
metadata: dict = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ConversationHistoryResponse(BaseModel):
|
||||
"""Response model for conversation history"""
|
||||
conversation_id: str
|
||||
turn_count: int
|
||||
total_tokens: int
|
||||
turns: List[ConversationTurnResponse]
|
||||
|
||||
|
||||
class SearchResultResponse(BaseModel):
|
||||
"""Response model for a single search result"""
|
||||
conversation_id: str
|
||||
turn_number: int
|
||||
role: str
|
||||
content: str
|
||||
timestamp: str
|
||||
score: float
|
||||
|
||||
|
||||
class SearchResponse(BaseModel):
|
||||
"""Response model for search results"""
|
||||
query: str
|
||||
results: List[SearchResultResponse]
|
||||
count: int
|
||||
|
||||
|
||||
class ConversationStatsResponse(BaseModel):
|
||||
"""Response model for conversation statistics"""
|
||||
conversation_id: str
|
||||
buffer_turns: int
|
||||
buffer_tokens: int
|
||||
qdrant_turns: int
|
||||
qdrant_tokens: int
|
||||
exists_in_buffer: bool
|
||||
exists_in_qdrant: bool
|
||||
|
||||
|
||||
class DeleteResponse(BaseModel):
|
||||
"""Response model for delete operation"""
|
||||
conversation_id: str
|
||||
deleted: bool
|
||||
message: str
|
||||
|
||||
|
||||
# Endpoints
|
||||
|
||||
@router.get(
|
||||
"/{conversation_id}",
|
||||
response_model=ConversationHistoryResponse,
|
||||
summary="Get conversation history",
|
||||
description="Retrieve complete conversation history including all turns"
|
||||
)
|
||||
async def get_conversation(
|
||||
conversation_id: str,
|
||||
include_buffer: bool = Query(
|
||||
True,
|
||||
description="Include recent turns from buffer that haven't been consolidated yet"
|
||||
)
|
||||
):
|
||||
"""
|
||||
Get complete conversation history
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
include_buffer: Include recent buffer turns not yet consolidated
|
||||
|
||||
Returns:
|
||||
Complete conversation history with all turns
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Get full history
|
||||
turns = await manager.get_full_history(conversation_id, include_buffer=include_buffer)
|
||||
|
||||
if not turns:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Conversation {conversation_id} not found"
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
turn_responses = []
|
||||
total_tokens = 0
|
||||
|
||||
for turn in turns:
|
||||
turn_response = ConversationTurnResponse(
|
||||
turn_number=turn.turn_number,
|
||||
role=turn.role.value if isinstance(turn.role, MessageRole) else turn.role,
|
||||
content=turn.content,
|
||||
timestamp=turn.timestamp.isoformat(),
|
||||
metadata=turn.metadata
|
||||
)
|
||||
|
||||
if turn.tokens:
|
||||
turn_response.tokens_prompt = turn.tokens.prompt
|
||||
turn_response.tokens_completion = turn.tokens.completion
|
||||
turn_response.tokens_total = turn.tokens.total
|
||||
total_tokens += turn.tokens.total
|
||||
|
||||
turn_responses.append(turn_response)
|
||||
|
||||
return ConversationHistoryResponse(
|
||||
conversation_id=conversation_id,
|
||||
turn_count=len(turns),
|
||||
total_tokens=total_tokens,
|
||||
turns=turn_responses
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{conversation_id}/stats",
|
||||
response_model=ConversationStatsResponse,
|
||||
summary="Get conversation statistics",
|
||||
description="Get detailed statistics about a conversation across all storage tiers"
|
||||
)
|
||||
async def get_conversation_stats(conversation_id: str):
|
||||
"""
|
||||
Get conversation statistics
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Statistics including turn counts and token usage across tiers
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
stats = await manager.get_conversation_stats(conversation_id)
|
||||
|
||||
return ConversationStatsResponse(**stats)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{conversation_id}/search",
|
||||
response_model=SearchResponse,
|
||||
summary="Search conversation semantically",
|
||||
description="Search for relevant turns within a conversation using semantic similarity"
|
||||
)
|
||||
async def search_conversation(
|
||||
conversation_id: str,
|
||||
search_request: SearchRequest
|
||||
):
|
||||
"""
|
||||
Semantic search within a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
search_request: Search query and parameters
|
||||
|
||||
Returns:
|
||||
Relevant conversation turns ranked by semantic similarity
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Perform semantic search
|
||||
results = await manager.search_conversations(
|
||||
query=search_request.query,
|
||||
conversation_id=conversation_id,
|
||||
limit=search_request.limit
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
search_results = [
|
||||
SearchResultResponse(
|
||||
conversation_id=result["conversation_id"],
|
||||
turn_number=result["turn_number"],
|
||||
role=result["role"],
|
||||
content=result["content"],
|
||||
timestamp=result["timestamp"],
|
||||
score=result["score"]
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
return SearchResponse(
|
||||
query=search_request.query,
|
||||
results=search_results,
|
||||
count=len(search_results)
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/search",
|
||||
response_model=SearchResponse,
|
||||
summary="Search all conversations",
|
||||
description="Search across all conversations using semantic similarity"
|
||||
)
|
||||
async def search_all_conversations(search_request: SearchRequest):
|
||||
"""
|
||||
Semantic search across all conversations
|
||||
|
||||
Args:
|
||||
search_request: Search query and parameters
|
||||
|
||||
Returns:
|
||||
Relevant turns from any conversation ranked by semantic similarity
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Perform semantic search across all conversations
|
||||
results = await manager.search_conversations(
|
||||
query=search_request.query,
|
||||
conversation_id=None, # Search all conversations
|
||||
limit=search_request.limit
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
search_results = [
|
||||
SearchResultResponse(
|
||||
conversation_id=result["conversation_id"],
|
||||
turn_number=result["turn_number"],
|
||||
role=result["role"],
|
||||
content=result["content"],
|
||||
timestamp=result["timestamp"],
|
||||
score=result["score"]
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
return SearchResponse(
|
||||
query=search_request.query,
|
||||
results=search_results,
|
||||
count=len(search_results)
|
||||
)
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/{conversation_id}",
|
||||
response_model=DeleteResponse,
|
||||
summary="Delete conversation",
|
||||
description="Delete a conversation from all storage tiers"
|
||||
)
|
||||
async def delete_conversation(
|
||||
conversation_id: str,
|
||||
clear_buffer: bool = Query(True, description="Clear from buffer (Tier 1)"),
|
||||
clear_qdrant: bool = Query(True, description="Clear from Qdrant (Tier 2/3)")
|
||||
):
|
||||
"""
|
||||
Delete a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
clear_buffer: Clear from Tier 1 buffer
|
||||
clear_qdrant: Clear from Tier 2/3 Qdrant
|
||||
|
||||
Returns:
|
||||
Deletion confirmation
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
try:
|
||||
await manager.clear_conversation(
|
||||
conversation_id,
|
||||
clear_buffer=clear_buffer,
|
||||
clear_qdrant=clear_qdrant
|
||||
)
|
||||
|
||||
return DeleteResponse(
|
||||
conversation_id=conversation_id,
|
||||
deleted=True,
|
||||
message=f"Conversation {conversation_id} deleted successfully"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error deleting conversation: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{conversation_id}/consolidate",
|
||||
summary="Consolidate conversation",
|
||||
description="Manually trigger consolidation from buffer to persistent storage"
|
||||
)
|
||||
async def consolidate_conversation(conversation_id: str):
|
||||
"""
|
||||
Manually consolidate a conversation
|
||||
|
||||
Moves all buffer turns to Qdrant persistent storage.
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Number of turns consolidated
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
try:
|
||||
count = await manager.consolidate(conversation_id)
|
||||
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
"consolidated_turns": count,
|
||||
"message": f"Successfully consolidated {count} turns to persistent storage"
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error consolidating conversation: {str(e)}"
|
||||
)
|
||||
@@ -1,34 +0,0 @@
|
||||
"""
|
||||
OpenAI-compatible /v1/models endpoint
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter
|
||||
from .schemas import ModelsListResponse, ModelInfo
|
||||
from src.config import get_settings
|
||||
|
||||
router = APIRouter()
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
@router.get("/v1/models")
|
||||
async def list_models():
|
||||
"""List available models in OpenAI format."""
|
||||
|
||||
models = []
|
||||
|
||||
# Add OpenAI-style aliases
|
||||
for alias in settings.model_aliases.keys():
|
||||
models.append(ModelInfo(id=alias, owned_by="tatlock"))
|
||||
|
||||
# Add actual local models
|
||||
for model_list in [
|
||||
settings.get_lightweight_models(),
|
||||
settings.get_heavy_models(),
|
||||
settings.get_code_models()
|
||||
]:
|
||||
for model in model_list:
|
||||
# Avoid duplicates
|
||||
if model not in [m.id for m in models]:
|
||||
models.append(ModelInfo(id=model, owned_by="tatlock"))
|
||||
|
||||
return ModelsListResponse(data=models)
|
||||
@@ -1,154 +0,0 @@
|
||||
"""
|
||||
OpenAI-compatible API schemas for /v1/* endpoints
|
||||
"""
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Literal
|
||||
from enum import Enum
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Request Schemas
|
||||
# ============================================================================
|
||||
|
||||
class MessageRole(str, Enum):
|
||||
"""Valid message roles."""
|
||||
SYSTEM = "system"
|
||||
USER = "user"
|
||||
ASSISTANT = "assistant"
|
||||
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
"""A single message in the conversation."""
|
||||
role: MessageRole
|
||||
content: str
|
||||
name: Optional[str] = None
|
||||
|
||||
|
||||
class ChatCompletionRequest(BaseModel):
|
||||
"""OpenAI-compatible chat completion request."""
|
||||
|
||||
model: str = Field(..., description="Model to use")
|
||||
messages: List[ChatMessage] = Field(..., min_length=1)
|
||||
stream: bool = Field(default=False, description="Enable streaming")
|
||||
|
||||
# Memory system (Phase 2)
|
||||
conversation_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Conversation ID for memory tracking (auto-generated if not provided)"
|
||||
)
|
||||
store_in_memory: bool = Field(
|
||||
default=True,
|
||||
description="Store conversation turns in memory system"
|
||||
)
|
||||
|
||||
# Multi-tenancy (Phase 2.5)
|
||||
user_id: str = Field(
|
||||
default="llm-testuser",
|
||||
description="User ID for multi-tenant memory isolation (future: extracted from auth token)"
|
||||
)
|
||||
|
||||
# A/B Testing (Phase 3.5)
|
||||
system_prompt_override: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Override system prompt variant for A/B testing (v1_verbose, v2_concise, v3_imperative, v4_minimal, v4_gemini_suggestion)"
|
||||
)
|
||||
|
||||
# Optional parameters
|
||||
temperature: Optional[float] = Field(default=0.7, ge=0, le=2)
|
||||
top_p: Optional[float] = Field(default=1.0, ge=0, le=1)
|
||||
max_tokens: Optional[int] = Field(default=None, ge=1)
|
||||
frequency_penalty: Optional[float] = Field(default=0.0, ge=-2, le=2)
|
||||
presence_penalty: Optional[float] = Field(default=0.0, ge=-2, le=2)
|
||||
stop: Optional[List[str]] = None
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello!"}
|
||||
],
|
||||
"stream": False,
|
||||
"temperature": 0.7
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Response Schemas
|
||||
# ============================================================================
|
||||
|
||||
class ChatMessageResponse(BaseModel):
|
||||
"""Response message."""
|
||||
role: str = "assistant"
|
||||
content: str
|
||||
|
||||
|
||||
class ChatCompletionChoice(BaseModel):
|
||||
"""A single completion choice."""
|
||||
index: int = 0
|
||||
message: ChatMessageResponse
|
||||
finish_reason: str = "stop"
|
||||
|
||||
|
||||
class UsageInfo(BaseModel):
|
||||
"""Token usage information."""
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
|
||||
|
||||
class ChatCompletionResponse(BaseModel):
|
||||
"""OpenAI-compatible chat completion response (non-streaming)."""
|
||||
|
||||
id: str
|
||||
object: str = "chat.completion"
|
||||
created: int
|
||||
model: str
|
||||
choices: List[ChatCompletionChoice]
|
||||
usage: UsageInfo
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Streaming Response Schemas
|
||||
# ============================================================================
|
||||
|
||||
class DeltaMessage(BaseModel):
|
||||
"""Delta message for streaming."""
|
||||
role: Optional[str] = None
|
||||
content: Optional[str] = None
|
||||
|
||||
|
||||
class ChatCompletionStreamChoice(BaseModel):
|
||||
"""Streaming choice."""
|
||||
index: int = 0
|
||||
delta: DeltaMessage
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
|
||||
class ChatCompletionStreamResponse(BaseModel):
|
||||
"""OpenAI-compatible streaming chunk."""
|
||||
|
||||
id: str
|
||||
object: str = "chat.completion.chunk"
|
||||
created: int
|
||||
model: str
|
||||
choices: List[ChatCompletionStreamChoice]
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Models Endpoint
|
||||
# ============================================================================
|
||||
|
||||
class ModelInfo(BaseModel):
|
||||
"""Model information."""
|
||||
id: str
|
||||
object: str = "model"
|
||||
created: int = 0
|
||||
owned_by: str = "local"
|
||||
|
||||
|
||||
class ModelsListResponse(BaseModel):
|
||||
"""List of available models."""
|
||||
object: str = "list"
|
||||
data: List[ModelInfo]
|
||||
@@ -1,858 +0,0 @@
|
||||
"""
|
||||
AI Controller
|
||||
|
||||
Provides AI orchestration endpoints including:
|
||||
- OpenAI-compatible chat completions
|
||||
- Model listing
|
||||
- Conversation memory management
|
||||
"""
|
||||
import time
|
||||
import logging
|
||||
import uuid
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from fastapi.responses import StreamingResponse
|
||||
from typing import AsyncIterator, List, Optional
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from src.controllers.base import BaseController
|
||||
from src.api.v1.schemas import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionChoice,
|
||||
ChatMessageResponse,
|
||||
UsageInfo,
|
||||
ChatCompletionStreamResponse,
|
||||
ChatCompletionStreamChoice,
|
||||
DeltaMessage,
|
||||
ModelsListResponse,
|
||||
ModelInfo,
|
||||
)
|
||||
from src.models.ollama_client import get_ollama_client
|
||||
from src.memory import get_memory_manager, MessageRole as MemoryMessageRole, TokenUsage
|
||||
from src.config import get_settings
|
||||
|
||||
# Agent orchestration
|
||||
try:
|
||||
from src.agent import get_unified_agent
|
||||
from src.agent.streaming import stream_agent_to_sse
|
||||
AGENT_AVAILABLE = True
|
||||
except ImportError as e:
|
||||
AGENT_AVAILABLE = False
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(f"Agent not available: {e}")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Request/Response Models for Conversations
|
||||
class SearchRequest(BaseModel):
|
||||
"""Request model for semantic search"""
|
||||
query: str = Field(..., description="Search query")
|
||||
limit: int = Field(5, ge=1, le=50, description="Maximum number of results")
|
||||
|
||||
|
||||
class ConversationTurnResponse(BaseModel):
|
||||
"""Response model for a conversation turn"""
|
||||
turn_number: int
|
||||
role: str
|
||||
content: str
|
||||
timestamp: str
|
||||
tokens_prompt: Optional[int] = None
|
||||
tokens_completion: Optional[int] = None
|
||||
tokens_total: Optional[int] = None
|
||||
metadata: dict = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ConversationHistoryResponse(BaseModel):
|
||||
"""Response model for conversation history"""
|
||||
conversation_id: str
|
||||
turn_count: int
|
||||
total_tokens: int
|
||||
turns: List[ConversationTurnResponse]
|
||||
|
||||
|
||||
class SearchResultResponse(BaseModel):
|
||||
"""Response model for a single search result"""
|
||||
conversation_id: str
|
||||
turn_number: int
|
||||
role: str
|
||||
content: str
|
||||
timestamp: str
|
||||
score: float
|
||||
|
||||
|
||||
class SearchResponse(BaseModel):
|
||||
"""Response model for search results"""
|
||||
query: str
|
||||
results: List[SearchResultResponse]
|
||||
count: int
|
||||
|
||||
|
||||
class ConversationStatsResponse(BaseModel):
|
||||
"""Response model for conversation statistics"""
|
||||
conversation_id: str
|
||||
buffer_turns: int
|
||||
buffer_tokens: int
|
||||
qdrant_turns: int
|
||||
qdrant_tokens: int
|
||||
exists_in_buffer: bool
|
||||
exists_in_qdrant: bool
|
||||
|
||||
|
||||
class DeleteResponse(BaseModel):
|
||||
"""Response model for delete operation"""
|
||||
conversation_id: str
|
||||
deleted: bool
|
||||
message: str
|
||||
|
||||
|
||||
# Helper functions
|
||||
def build_prompt_from_messages(messages: list) -> str:
|
||||
"""
|
||||
Convert message list to a prompt string.
|
||||
"""
|
||||
prompt_parts = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.role.value if hasattr(msg.role, 'value') else msg.role
|
||||
content = msg.content
|
||||
|
||||
if role == "system":
|
||||
prompt_parts.append(f"System: {content}")
|
||||
elif role == "user":
|
||||
prompt_parts.append(f"User: {content}")
|
||||
elif role == "assistant":
|
||||
prompt_parts.append(f"Assistant: {content}")
|
||||
|
||||
prompt_parts.append("Assistant:")
|
||||
return "\n\n".join(prompt_parts)
|
||||
|
||||
|
||||
async def stream_chat_completion(
|
||||
request_id: str,
|
||||
model: str,
|
||||
prompt: str,
|
||||
temperature: float,
|
||||
max_tokens: int | None
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream chat completion in OpenAI SSE format.
|
||||
|
||||
Yields:
|
||||
Server-Sent Events formatted strings
|
||||
"""
|
||||
created = int(time.time())
|
||||
ollama_client = get_ollama_client()
|
||||
|
||||
# First chunk with role
|
||||
first_chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(role="assistant"),
|
||||
finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {first_chunk.model_dump_json()}\n\n"
|
||||
|
||||
# Stream tokens
|
||||
try:
|
||||
async for token in ollama_client.generate_streaming(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens
|
||||
):
|
||||
chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(content=token),
|
||||
finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Streaming error: {e}")
|
||||
# Send error in OpenAI format
|
||||
import json
|
||||
error_chunk = {
|
||||
"error": {
|
||||
"message": str(e),
|
||||
"type": "server_error"
|
||||
}
|
||||
}
|
||||
yield f"data: {json.dumps(error_chunk)}\n\n"
|
||||
return
|
||||
|
||||
# Final chunk
|
||||
final_chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(),
|
||||
finish_reason="stop"
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {final_chunk.model_dump_json()}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
|
||||
async def store_conversation_turn(
|
||||
conversation_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
user_id: str = "llm-testuser",
|
||||
tokens: dict = None
|
||||
):
|
||||
"""
|
||||
Store a conversation turn in memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
role: Message role (user, assistant, system)
|
||||
content: Message content
|
||||
user_id: User ID for multi-tenancy (defaults to "llm-testuser")
|
||||
tokens: Optional token usage dict
|
||||
"""
|
||||
try:
|
||||
# Don't store system messages - they're part of the agent's state_modifier
|
||||
if role == "system":
|
||||
logger.debug(f"Skipping storage of system message for {conversation_id}")
|
||||
return
|
||||
|
||||
memory_manager = get_memory_manager()
|
||||
|
||||
# Convert role string to MemoryMessageRole
|
||||
if role == "user":
|
||||
memory_role = MemoryMessageRole.USER
|
||||
elif role == "assistant":
|
||||
memory_role = MemoryMessageRole.ASSISTANT
|
||||
else:
|
||||
memory_role = MemoryMessageRole.USER # Default fallback
|
||||
|
||||
# Create TokenUsage if provided
|
||||
token_usage = None
|
||||
if tokens:
|
||||
token_usage = TokenUsage(
|
||||
prompt=tokens.get("prompt", 0),
|
||||
completion=tokens.get("completion", 0),
|
||||
total=tokens.get("total", 0)
|
||||
)
|
||||
|
||||
# Store in memory with user_id
|
||||
await memory_manager.add_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=memory_role,
|
||||
content=content,
|
||||
user_id=user_id,
|
||||
tokens=token_usage
|
||||
)
|
||||
|
||||
logger.debug(f"Stored {role} turn in memory for user={user_id}, conversation={conversation_id}")
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the request
|
||||
logger.error(f"Failed to store turn in memory: {e}")
|
||||
|
||||
|
||||
class AIController(BaseController):
|
||||
"""
|
||||
Controller for AI orchestration operations
|
||||
|
||||
Provides endpoints for:
|
||||
- OpenAI-compatible chat completions (streaming and non-streaming)
|
||||
- Model listing
|
||||
- Conversation memory management
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(prefix="/v1", tags=["AI"])
|
||||
|
||||
def create_router(self) -> APIRouter:
|
||||
"""Create and configure the router"""
|
||||
router = APIRouter()
|
||||
settings = get_settings()
|
||||
|
||||
# Chat Completions Endpoint
|
||||
@router.post(
|
||||
"/v1/chat/completions",
|
||||
tags=["AI"]
|
||||
)
|
||||
async def chat_completions(request: ChatCompletionRequest):
|
||||
"""
|
||||
OpenAI-compatible chat completions endpoint.
|
||||
Supports both streaming and non-streaming.
|
||||
|
||||
Automatically stores conversations in memory system if enabled.
|
||||
"""
|
||||
request_id = f"chatcmpl-{int(time.time() * 1000)}"
|
||||
|
||||
# Multi-tenancy: Extract user_id (defaults to "llm-testuser")
|
||||
user_id = request.user_id
|
||||
|
||||
# Generate or use provided conversation_id
|
||||
conversation_id = request.conversation_id or f"conv_{uuid.uuid4().hex[:16]}"
|
||||
|
||||
logger.info(
|
||||
f"Chat request: id={request_id}, user={user_id}, model={request.model}, "
|
||||
f"messages={len(request.messages)}, stream={request.stream}, "
|
||||
f"conversation_id={conversation_id}, store_in_memory={request.store_in_memory}"
|
||||
)
|
||||
|
||||
# Always route through unified agent (with fallback to direct Ollama)
|
||||
if AGENT_AVAILABLE:
|
||||
try:
|
||||
logger.info(f"Using unified agent for request {request_id}")
|
||||
|
||||
# Get last message
|
||||
user_message = request.messages[-1].content
|
||||
|
||||
# Load conversation history from memory if available
|
||||
history = []
|
||||
if request.store_in_memory:
|
||||
memory_manager = get_memory_manager()
|
||||
|
||||
# Check both buffer (Tier 1) and Qdrant (Tier 2/3)
|
||||
buffer_exists = await memory_manager.buffer_memory.conversation_exists(conversation_id)
|
||||
qdrant_exists = await memory_manager.qdrant_memory.conversation_exists(conversation_id)
|
||||
conversation_exists = buffer_exists or qdrant_exists
|
||||
|
||||
if conversation_exists:
|
||||
# Load full history (combines buffer + Qdrant)
|
||||
logger.info(f"Loading conversation history from memory for {conversation_id} (buffer={buffer_exists}, qdrant={qdrant_exists})")
|
||||
all_turns = await memory_manager.get_full_history(conversation_id, include_buffer=True)
|
||||
# Get most recent 20 turns
|
||||
recent_turns = all_turns[-20:] if len(all_turns) > 20 else all_turns
|
||||
# Filter out system messages - they should not be in conversation history
|
||||
history = [
|
||||
{"role": turn.role.value, "content": turn.content}
|
||||
for turn in recent_turns
|
||||
if turn.role.value != "system"
|
||||
]
|
||||
logger.info(f"✓ Loaded {len(history)} turns from memory (total: {len(all_turns)})")
|
||||
else:
|
||||
# New conversation - use request messages (excluding system messages)
|
||||
logger.info(f"New conversation {conversation_id} - using request messages")
|
||||
for msg in request.messages[:-1]: # All except last
|
||||
if msg.role.value != "system": # Skip system messages
|
||||
history.append({"role": msg.role.value, "content": msg.content})
|
||||
else:
|
||||
# Memory disabled - fall back to request messages (excluding system messages)
|
||||
for msg in request.messages[:-1]:
|
||||
if msg.role.value != "system": # Skip system messages
|
||||
history.append({"role": msg.role.value, "content": msg.content})
|
||||
|
||||
# Store user message in memory BEFORE agent execution
|
||||
if request.store_in_memory:
|
||||
logger.info(f"Storing user message in memory for {conversation_id}")
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role="user",
|
||||
content=user_message,
|
||||
user_id=user_id
|
||||
)
|
||||
logger.info(f"✓ Stored user message in memory for {conversation_id}")
|
||||
|
||||
# Get agent
|
||||
agent = get_unified_agent()
|
||||
|
||||
# Stream response
|
||||
if request.stream:
|
||||
# For streaming, we need to collect the response to store it
|
||||
collected_content = []
|
||||
|
||||
async def agent_stream_generator():
|
||||
nonlocal collected_content
|
||||
try:
|
||||
agent_stream = agent.chat(
|
||||
message=user_message,
|
||||
conversation_history=history,
|
||||
stream=True,
|
||||
prompt_variant=request.system_prompt_override
|
||||
)
|
||||
# Always use "Tatlock" as model name in responses
|
||||
async for sse_chunk in stream_agent_to_sse(agent_stream, request_id, "Tatlock"):
|
||||
# Collect content for memory storage
|
||||
# Extract content from SSE chunk if it contains delta content
|
||||
if '"content":' in sse_chunk:
|
||||
try:
|
||||
import json
|
||||
# Parse the SSE data line
|
||||
for line in sse_chunk.split('\n'):
|
||||
if line.startswith('data: ') and not line.startswith('data: [DONE]'):
|
||||
chunk_data = json.loads(line[6:]) # Remove 'data: ' prefix
|
||||
if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
|
||||
delta = chunk_data['choices'][0].get('delta', {})
|
||||
if 'content' in delta:
|
||||
collected_content.append(delta['content'])
|
||||
except:
|
||||
pass
|
||||
yield sse_chunk.encode('utf-8')
|
||||
finally:
|
||||
# Store assistant response in memory AFTER streaming completes
|
||||
# This runs in the finally block to ensure it executes even if client disconnects
|
||||
if request.store_in_memory and collected_content:
|
||||
full_response = ''.join(collected_content)
|
||||
logger.info(f"Storing assistant response in memory for {conversation_id}")
|
||||
try:
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role="assistant",
|
||||
content=full_response,
|
||||
user_id=user_id
|
||||
)
|
||||
logger.info(f"✓ Stored assistant response in memory for {conversation_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to store assistant response: {e}")
|
||||
|
||||
return StreamingResponse(
|
||||
agent_stream_generator(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no"
|
||||
}
|
||||
)
|
||||
else:
|
||||
# Non-streaming
|
||||
response_text = await agent.chat_completion(
|
||||
message=user_message,
|
||||
conversation_history=history,
|
||||
prompt_variant=request.system_prompt_override
|
||||
)
|
||||
|
||||
# Store assistant response in memory AFTER agent execution
|
||||
if request.store_in_memory:
|
||||
# Estimate token usage (simple word count)
|
||||
prompt_tokens = len(user_message.split())
|
||||
completion_tokens = len(response_text.split())
|
||||
logger.info(f"Storing assistant response in memory for {conversation_id}")
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role="assistant",
|
||||
content=response_text,
|
||||
user_id=user_id,
|
||||
tokens={
|
||||
"prompt": prompt_tokens,
|
||||
"completion": completion_tokens,
|
||||
"total": prompt_tokens + completion_tokens
|
||||
}
|
||||
)
|
||||
logger.info(f"✓ Stored assistant response in memory for {conversation_id}")
|
||||
|
||||
# Always use "Tatlock" as model name in responses
|
||||
return ChatCompletionResponse(
|
||||
id=request_id,
|
||||
object="chat.completion",
|
||||
created=int(time.time()),
|
||||
model="Tatlock",
|
||||
choices=[
|
||||
ChatCompletionChoice(
|
||||
index=0,
|
||||
message=ChatMessageResponse(
|
||||
role="assistant",
|
||||
content=response_text
|
||||
),
|
||||
finish_reason="stop"
|
||||
)
|
||||
],
|
||||
usage=UsageInfo(
|
||||
prompt_tokens=len(user_message.split()),
|
||||
completion_tokens=len(response_text.split()),
|
||||
total_tokens=len(user_message.split()) + len(response_text.split())
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
if not settings.agent_fallback_enabled:
|
||||
logger.error(f"Agent failed and fallback is disabled. Error: {e}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Agent failed to generate completion: {str(e)}"
|
||||
)
|
||||
logger.error(f"Agent failed, falling back to direct Ollama: {e}")
|
||||
# Fall through to direct Ollama call below
|
||||
|
||||
# Store user messages in memory (if enabled)
|
||||
if request.store_in_memory:
|
||||
for msg in request.messages:
|
||||
role = msg.role.value if hasattr(msg.role, 'value') else msg.role
|
||||
if role == "user": # Store latest user message
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=role,
|
||||
content=msg.content,
|
||||
user_id=user_id
|
||||
)
|
||||
|
||||
# Build prompt from messages
|
||||
prompt = build_prompt_from_messages(request.messages)
|
||||
|
||||
# Streaming response
|
||||
if request.stream:
|
||||
return StreamingResponse(
|
||||
stream_chat_completion(
|
||||
request_id=request_id,
|
||||
model=request.model,
|
||||
prompt=prompt,
|
||||
temperature=request.temperature,
|
||||
max_tokens=request.max_tokens
|
||||
),
|
||||
media_type="text/event-stream"
|
||||
)
|
||||
|
||||
# Non-streaming response
|
||||
try:
|
||||
ollama_client = get_ollama_client()
|
||||
result = await ollama_client.generate_non_streaming(
|
||||
model=request.model,
|
||||
prompt=prompt,
|
||||
temperature=request.temperature,
|
||||
max_tokens=request.max_tokens
|
||||
)
|
||||
|
||||
assistant_content = result["response"]
|
||||
|
||||
# Store assistant response in memory (if enabled)
|
||||
if request.store_in_memory:
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role="assistant",
|
||||
content=assistant_content,
|
||||
user_id=user_id,
|
||||
tokens=result["tokens"]
|
||||
)
|
||||
|
||||
response = ChatCompletionResponse(
|
||||
id=request_id,
|
||||
created=int(time.time()),
|
||||
model=request.model,
|
||||
choices=[
|
||||
ChatCompletionChoice(
|
||||
index=0,
|
||||
message=ChatMessageResponse(
|
||||
role="assistant",
|
||||
content=assistant_content
|
||||
),
|
||||
finish_reason="stop"
|
||||
)
|
||||
],
|
||||
usage=UsageInfo(
|
||||
prompt_tokens=result["tokens"]["prompt"],
|
||||
completion_tokens=result["tokens"]["completion"],
|
||||
total_tokens=result["tokens"]["total"]
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Chat response: id={request_id}, "
|
||||
f"tokens={result['tokens']['total']}, "
|
||||
f"conversation_id={conversation_id}"
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Chat completion error: {e}")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to generate completion: {str(e)}"
|
||||
)
|
||||
|
||||
# Models Endpoint
|
||||
@router.get(
|
||||
"/v1/models",
|
||||
tags=["AI"]
|
||||
)
|
||||
async def list_models():
|
||||
"""List available models in OpenAI format."""
|
||||
# Unified agent - always uses mistral:7b with tools
|
||||
# Model name is "Tatlock" for all requests
|
||||
return ModelsListResponse(
|
||||
data=[
|
||||
ModelInfo(
|
||||
id="Tatlock",
|
||||
owned_by="tatlock",
|
||||
created=1640000000 # Fixed timestamp for consistency
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Conversation Endpoints
|
||||
@router.get(
|
||||
"/v1/conversations/{conversation_id}",
|
||||
response_model=ConversationHistoryResponse,
|
||||
tags=["Conversations"],
|
||||
summary="Get conversation history",
|
||||
description="Retrieve complete conversation history including all turns"
|
||||
)
|
||||
async def get_conversation(
|
||||
conversation_id: str,
|
||||
include_buffer: bool = Query(
|
||||
True,
|
||||
description="Include recent turns from buffer that haven't been consolidated yet"
|
||||
)
|
||||
):
|
||||
"""
|
||||
Get complete conversation history
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
include_buffer: Include recent buffer turns not yet consolidated
|
||||
|
||||
Returns:
|
||||
Complete conversation history with all turns
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Get full history
|
||||
turns = await manager.get_full_history(conversation_id, include_buffer=include_buffer)
|
||||
|
||||
if not turns:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Conversation {conversation_id} not found"
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
turn_responses = []
|
||||
total_tokens = 0
|
||||
|
||||
for turn in turns:
|
||||
turn_response = ConversationTurnResponse(
|
||||
turn_number=turn.turn_number,
|
||||
role=turn.role.value if isinstance(turn.role, MemoryMessageRole) else turn.role,
|
||||
content=turn.content,
|
||||
timestamp=turn.timestamp.isoformat(),
|
||||
metadata=turn.metadata
|
||||
)
|
||||
|
||||
if turn.tokens:
|
||||
turn_response.tokens_prompt = turn.tokens.prompt
|
||||
turn_response.tokens_completion = turn.tokens.completion
|
||||
turn_response.tokens_total = turn.tokens.total
|
||||
total_tokens += turn.tokens.total
|
||||
|
||||
turn_responses.append(turn_response)
|
||||
|
||||
return ConversationHistoryResponse(
|
||||
conversation_id=conversation_id,
|
||||
turn_count=len(turns),
|
||||
total_tokens=total_tokens,
|
||||
turns=turn_responses
|
||||
)
|
||||
|
||||
@router.get(
|
||||
"/v1/conversations/{conversation_id}/stats",
|
||||
response_model=ConversationStatsResponse,
|
||||
tags=["Conversations"],
|
||||
summary="Get conversation statistics",
|
||||
description="Get detailed statistics about a conversation across all storage tiers"
|
||||
)
|
||||
async def get_conversation_stats(conversation_id: str):
|
||||
"""
|
||||
Get conversation statistics
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Statistics including turn counts and token usage across tiers
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
stats = await manager.get_conversation_stats(conversation_id)
|
||||
|
||||
return ConversationStatsResponse(**stats)
|
||||
|
||||
@router.post(
|
||||
"/v1/conversations/{conversation_id}/search",
|
||||
response_model=SearchResponse,
|
||||
tags=["Conversations"],
|
||||
summary="Search conversation semantically",
|
||||
description="Search for relevant turns within a conversation using semantic similarity"
|
||||
)
|
||||
async def search_conversation(
|
||||
conversation_id: str,
|
||||
search_request: SearchRequest
|
||||
):
|
||||
"""
|
||||
Semantic search within a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
search_request: Search query and parameters
|
||||
|
||||
Returns:
|
||||
Relevant conversation turns ranked by semantic similarity
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Perform semantic search
|
||||
results = await manager.search_conversations(
|
||||
query=search_request.query,
|
||||
conversation_id=conversation_id,
|
||||
limit=search_request.limit
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
search_results = [
|
||||
SearchResultResponse(
|
||||
conversation_id=result["conversation_id"],
|
||||
turn_number=result["turn_number"],
|
||||
role=result["role"],
|
||||
content=result["content"],
|
||||
timestamp=result["timestamp"],
|
||||
score=result["score"]
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
return SearchResponse(
|
||||
query=search_request.query,
|
||||
results=search_results,
|
||||
count=len(search_results)
|
||||
)
|
||||
|
||||
@router.post(
|
||||
"/v1/conversations/search",
|
||||
response_model=SearchResponse,
|
||||
tags=["Conversations"],
|
||||
summary="Search all conversations",
|
||||
description="Search across all conversations using semantic similarity"
|
||||
)
|
||||
async def search_all_conversations(search_request: SearchRequest):
|
||||
"""
|
||||
Semantic search across all conversations
|
||||
|
||||
Args:
|
||||
search_request: Search query and parameters
|
||||
|
||||
Returns:
|
||||
Relevant turns from any conversation ranked by semantic similarity
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Perform semantic search across all conversations
|
||||
results = await manager.search_conversations(
|
||||
query=search_request.query,
|
||||
conversation_id=None, # Search all conversations
|
||||
limit=search_request.limit
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
search_results = [
|
||||
SearchResultResponse(
|
||||
conversation_id=result["conversation_id"],
|
||||
turn_number=result["turn_number"],
|
||||
role=result["role"],
|
||||
content=result["content"],
|
||||
timestamp=result["timestamp"],
|
||||
score=result["score"]
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
return SearchResponse(
|
||||
query=search_request.query,
|
||||
results=search_results,
|
||||
count=len(search_results)
|
||||
)
|
||||
|
||||
@router.delete(
|
||||
"/v1/conversations/{conversation_id}",
|
||||
response_model=DeleteResponse,
|
||||
tags=["Conversations"],
|
||||
summary="Delete conversation",
|
||||
description="Delete a conversation from all storage tiers"
|
||||
)
|
||||
async def delete_conversation(
|
||||
conversation_id: str,
|
||||
clear_buffer: bool = Query(True, description="Clear from buffer (Tier 1)"),
|
||||
clear_qdrant: bool = Query(True, description="Clear from Qdrant (Tier 2/3)")
|
||||
):
|
||||
"""
|
||||
Delete a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
clear_buffer: Clear from Tier 1 buffer
|
||||
clear_qdrant: Clear from Tier 2/3 Qdrant
|
||||
|
||||
Returns:
|
||||
Deletion confirmation
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
try:
|
||||
await manager.clear_conversation(
|
||||
conversation_id,
|
||||
clear_buffer=clear_buffer,
|
||||
clear_qdrant=clear_qdrant
|
||||
)
|
||||
|
||||
return DeleteResponse(
|
||||
conversation_id=conversation_id,
|
||||
deleted=True,
|
||||
message=f"Conversation {conversation_id} deleted successfully"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error deleting conversation: {str(e)}"
|
||||
)
|
||||
|
||||
@router.post(
|
||||
"/v1/conversations/{conversation_id}/consolidate",
|
||||
tags=["Conversations"],
|
||||
summary="Consolidate conversation",
|
||||
description="Manually trigger consolidation from buffer to persistent storage"
|
||||
)
|
||||
async def consolidate_conversation(conversation_id: str):
|
||||
"""
|
||||
Manually consolidate a conversation
|
||||
|
||||
Moves all buffer turns to Qdrant persistent storage.
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Number of turns consolidated
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
try:
|
||||
count = await manager.consolidate(conversation_id)
|
||||
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
"consolidated_turns": count,
|
||||
"message": f"Successfully consolidated {count} turns to persistent storage"
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error consolidating conversation: {str(e)}"
|
||||
)
|
||||
|
||||
return router
|
||||
|
||||
|
||||
# Create controller instance
|
||||
ai_controller = AIController()
|
||||
@@ -10,7 +10,6 @@ from src.config import get_settings
|
||||
from src.logging_config import setup_logging, get_logger
|
||||
from src.models.ollama_client import get_ollama_client, close_ollama_client
|
||||
from src.controllers.infrastructure_controller import infrastructure_controller
|
||||
from src.controllers.ai_controller import ai_controller
|
||||
from src.controllers.tools_controller import tools_controller
|
||||
from src.controllers.health_controller import health_controller
|
||||
from src.controllers.static_controller import static_controller
|
||||
@@ -148,7 +147,6 @@ app.add_middleware(
|
||||
|
||||
# Include controller routers
|
||||
app.include_router(health_controller.router) # / and /health
|
||||
app.include_router(ai_controller.router) # /v1/chat, /v1/models, /v1/conversations
|
||||
app.include_router(tools_controller.router) # /web-scraper/scrape
|
||||
app.include_router(infrastructure_controller.router) # /infrastructure/*
|
||||
app.include_router(static_controller.router) # /static/*
|
||||
|
||||
@@ -1,42 +0,0 @@
|
||||
"""
|
||||
Memory system for conversation persistence
|
||||
|
||||
Simplified architecture:
|
||||
- Tier 1: ConversationBufferMemory (in-memory, fast, last 10 turns)
|
||||
- Tier 2/3: QdrantConversationMemory (unified persistent + semantic search)
|
||||
- Manager: MemoryManager (orchestrates all tiers)
|
||||
"""
|
||||
from .tier1_buffer import ConversationBufferMemory, get_buffer_memory
|
||||
from .qdrant_memory import QdrantConversationMemory, get_qdrant_memory
|
||||
from .manager import MemoryManager, get_memory_manager
|
||||
from .schemas import (
|
||||
ConversationTurn,
|
||||
ConversationBuffer,
|
||||
ConversationMetadata,
|
||||
ConversationSummary,
|
||||
MemoryQuery,
|
||||
MemoryResult,
|
||||
MessageRole,
|
||||
TokenUsage
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Manager (primary interface)
|
||||
"MemoryManager",
|
||||
"get_memory_manager",
|
||||
# Tier 1
|
||||
"ConversationBufferMemory",
|
||||
"get_buffer_memory",
|
||||
# Tier 2/3
|
||||
"QdrantConversationMemory",
|
||||
"get_qdrant_memory",
|
||||
# Schemas
|
||||
"ConversationTurn",
|
||||
"ConversationBuffer",
|
||||
"ConversationMetadata",
|
||||
"ConversationSummary",
|
||||
"MemoryQuery",
|
||||
"MemoryResult",
|
||||
"MessageRole",
|
||||
"TokenUsage",
|
||||
]
|
||||
@@ -1,169 +0,0 @@
|
||||
"""
|
||||
Base classes for memory system
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Optional
|
||||
from .schemas import ConversationTurn, ConversationBuffer, MemoryQuery, MemoryResult
|
||||
|
||||
|
||||
class BaseMemory(ABC):
|
||||
"""Base class for all memory tiers"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a new turn to memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns from memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists in this memory tier
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation exists
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier1Memory(BaseMemory):
|
||||
"""Base class for Tier 1 (working memory)"""
|
||||
|
||||
@abstractmethod
|
||||
async def get_buffer(self, conversation_id: str) -> Optional[ConversationBuffer]:
|
||||
"""
|
||||
Get the full conversation buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
ConversationBuffer or None if not found
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def prune(self, conversation_id: str, keep_last: int = 5) -> None:
|
||||
"""
|
||||
Prune old turns, keeping only the most recent ones
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
keep_last: Number of recent turns to keep
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier2Memory(BaseMemory):
|
||||
"""Base class for Tier 2 (short-term memory with summaries)"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_summary(
|
||||
self,
|
||||
conversation_id: str,
|
||||
summary_text: str,
|
||||
turn_range_start: int,
|
||||
turn_range_end: int
|
||||
) -> None:
|
||||
"""
|
||||
Add a conversation summary
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
summary_text: The summarized text
|
||||
turn_range_start: First turn number in summary
|
||||
turn_range_end: Last turn number in summary
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_summaries(self, conversation_id: str) -> List[dict]:
|
||||
"""
|
||||
Get all summaries for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
List of summary dictionaries
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier3Memory(BaseMemory):
|
||||
"""Base class for Tier 3 (long-term vector memory)"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_turn_with_embedding(
|
||||
self,
|
||||
conversation_id: str,
|
||||
turn: ConversationTurn,
|
||||
embedding: List[float]
|
||||
) -> None:
|
||||
"""
|
||||
Add a turn with its vector embedding
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn
|
||||
embedding: Vector embedding of the turn content
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def similarity_search(
|
||||
self,
|
||||
query_embedding: List[float],
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[dict]:
|
||||
"""
|
||||
Perform semantic similarity search
|
||||
|
||||
Args:
|
||||
query_embedding: Vector embedding of the search query
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
pass
|
||||
@@ -1,322 +0,0 @@
|
||||
"""
|
||||
Memory Manager: Orchestrates all memory tiers
|
||||
|
||||
Coordinates:
|
||||
- Tier 1: ConversationBufferMemory (RAM, fast, last N turns)
|
||||
- Tier 2/3: QdrantConversationMemory (persistent + semantic)
|
||||
|
||||
Provides unified interface for memory operations with automatic
|
||||
tier management and consolidation.
|
||||
"""
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
from .tier1_buffer import ConversationBufferMemory, get_buffer_memory
|
||||
from .qdrant_memory import QdrantConversationMemory, get_qdrant_memory
|
||||
from .schemas import ConversationTurn, MessageRole, TokenUsage
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class MemoryManager:
|
||||
"""
|
||||
Unified memory manager orchestrating all tiers
|
||||
|
||||
Responsibilities:
|
||||
- Add turns to appropriate tiers
|
||||
- Retrieve conversation history (buffer + persistent)
|
||||
- Consolidate buffer to persistent storage
|
||||
- Semantic search across all conversations
|
||||
- Memory lifecycle management
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
buffer_memory: Optional[ConversationBufferMemory] = None,
|
||||
qdrant_memory: Optional[QdrantConversationMemory] = None,
|
||||
auto_consolidate: bool = True
|
||||
):
|
||||
"""
|
||||
Initialize memory manager
|
||||
|
||||
Args:
|
||||
buffer_memory: Optional Tier 1 buffer instance
|
||||
qdrant_memory: Optional Tier 2/3 Qdrant instance
|
||||
auto_consolidate: Automatically consolidate when buffer threshold reached
|
||||
"""
|
||||
self.buffer_memory = buffer_memory or get_buffer_memory()
|
||||
self.qdrant_memory = qdrant_memory or get_qdrant_memory()
|
||||
self.auto_consolidate = auto_consolidate
|
||||
|
||||
logger.info(
|
||||
f"MemoryManager initialized (auto_consolidate={auto_consolidate})"
|
||||
)
|
||||
|
||||
async def add_turn(
|
||||
self,
|
||||
conversation_id: str,
|
||||
role: MessageRole,
|
||||
content: str,
|
||||
user_id: str = "llm-testuser",
|
||||
tokens: Optional[TokenUsage] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> ConversationTurn:
|
||||
"""
|
||||
Add a conversation turn to memory
|
||||
|
||||
Automatically:
|
||||
1. Adds to Tier 1 (buffer)
|
||||
2. Checks if consolidation threshold reached
|
||||
3. Consolidates to Tier 2/3 if needed
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
role: Message role (user, assistant, system)
|
||||
content: Message content
|
||||
user_id: User ID for multi-tenancy (defaults to "llm-testuser")
|
||||
tokens: Optional token usage
|
||||
metadata: Optional metadata
|
||||
|
||||
Returns:
|
||||
The created conversation turn
|
||||
"""
|
||||
# Get current buffer to determine turn number
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
turn_number = (buffer.metadata.turn_count + 1) if buffer else 1
|
||||
|
||||
# Create turn with user_id
|
||||
turn = ConversationTurn(
|
||||
role=role,
|
||||
content=content,
|
||||
timestamp=datetime.utcnow(),
|
||||
turn_number=turn_number,
|
||||
user_id=user_id,
|
||||
tokens=tokens,
|
||||
metadata=metadata or {}
|
||||
)
|
||||
|
||||
# Add to Tier 1 (buffer)
|
||||
await self.buffer_memory.add_turn(conversation_id, turn)
|
||||
logger.debug(f"Turn {turn_number} added to buffer for {conversation_id}")
|
||||
|
||||
# Check consolidation threshold
|
||||
if self.auto_consolidate:
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
if buffer.metadata.turn_count >= settings.memory_consolidation_threshold:
|
||||
logger.info(
|
||||
f"Consolidation threshold reached for {conversation_id} "
|
||||
f"({buffer.metadata.turn_count} turns)"
|
||||
)
|
||||
await self._consolidate_buffer(conversation_id)
|
||||
|
||||
return turn
|
||||
|
||||
async def get_recent_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 10
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get recent conversation turns (from buffer)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
|
||||
Returns:
|
||||
List of recent conversation turns
|
||||
"""
|
||||
return await self.buffer_memory.get_recent_turns(conversation_id, limit)
|
||||
|
||||
async def get_full_history(
|
||||
self,
|
||||
conversation_id: str,
|
||||
include_buffer: bool = True
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get complete conversation history
|
||||
|
||||
Combines:
|
||||
- Tier 2/3: Persistent history from Qdrant
|
||||
- Tier 1: Recent buffer (if include_buffer=True)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
include_buffer: Include buffer turns not yet consolidated
|
||||
|
||||
Returns:
|
||||
Complete conversation history, sorted chronologically
|
||||
"""
|
||||
# Get from Qdrant (Tier 2)
|
||||
qdrant_turns = await self.qdrant_memory.get_turns(conversation_id)
|
||||
|
||||
# Get from buffer (Tier 1)
|
||||
if include_buffer:
|
||||
buffer_turns = await self.buffer_memory.get_turns(conversation_id)
|
||||
|
||||
# Combine and deduplicate (Qdrant is source of truth)
|
||||
qdrant_turn_numbers = {t.turn_number for t in qdrant_turns}
|
||||
new_buffer_turns = [
|
||||
t for t in buffer_turns
|
||||
if t.turn_number not in qdrant_turn_numbers
|
||||
]
|
||||
|
||||
all_turns = qdrant_turns + new_buffer_turns
|
||||
else:
|
||||
all_turns = qdrant_turns
|
||||
|
||||
# Sort chronologically
|
||||
all_turns.sort(key=lambda t: t.turn_number)
|
||||
|
||||
return all_turns
|
||||
|
||||
async def search_conversations(
|
||||
self,
|
||||
query: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Semantic search across conversations (Tier 3 mode)
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
return await self.qdrant_memory.similarity_search(
|
||||
query=query,
|
||||
conversation_id=conversation_id,
|
||||
limit=limit
|
||||
)
|
||||
|
||||
async def consolidate(self, conversation_id: str) -> int:
|
||||
"""
|
||||
Manually trigger consolidation for a conversation
|
||||
|
||||
Moves all buffer turns to Qdrant (Tier 1 → Tier 2/3)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Number of turns consolidated
|
||||
"""
|
||||
return await self._consolidate_buffer(conversation_id)
|
||||
|
||||
async def _consolidate_buffer(self, conversation_id: str) -> int:
|
||||
"""
|
||||
Internal consolidation: Move buffer turns to Qdrant
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Number of turns consolidated
|
||||
"""
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
if not buffer or len(buffer.turns) == 0:
|
||||
logger.debug(f"No turns to consolidate for {conversation_id}")
|
||||
return 0
|
||||
|
||||
# Get turns from buffer
|
||||
buffer_turns = buffer.turns.copy()
|
||||
|
||||
# Add to Qdrant
|
||||
consolidated_count = 0
|
||||
for turn in buffer_turns:
|
||||
try:
|
||||
await self.qdrant_memory.add_turn(conversation_id, turn)
|
||||
consolidated_count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Error consolidating turn {turn.turn_number}: {e}")
|
||||
|
||||
logger.info(
|
||||
f"Consolidated {consolidated_count}/{len(buffer_turns)} turns "
|
||||
f"for {conversation_id}"
|
||||
)
|
||||
|
||||
# Note: We keep the buffer, just stored in Qdrant as well
|
||||
# Buffer will be pruned naturally as new turns come in
|
||||
# This provides redundancy and fast access to recent turns
|
||||
|
||||
return consolidated_count
|
||||
|
||||
async def clear_conversation(
|
||||
self,
|
||||
conversation_id: str,
|
||||
clear_buffer: bool = True,
|
||||
clear_qdrant: bool = True
|
||||
) -> None:
|
||||
"""
|
||||
Clear conversation from memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
clear_buffer: Clear from Tier 1 buffer
|
||||
clear_qdrant: Clear from Tier 2/3 Qdrant
|
||||
"""
|
||||
if clear_buffer:
|
||||
await self.buffer_memory.clear_conversation(conversation_id)
|
||||
logger.info(f"Cleared buffer for {conversation_id}")
|
||||
|
||||
if clear_qdrant:
|
||||
await self.qdrant_memory.clear_conversation(conversation_id)
|
||||
logger.info(f"Cleared Qdrant for {conversation_id}")
|
||||
|
||||
async def get_conversation_stats(
|
||||
self,
|
||||
conversation_id: str
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Get conversation statistics across all tiers
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Dictionary with stats from buffer and Qdrant
|
||||
"""
|
||||
# Get buffer stats
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
buffer_stats = {
|
||||
"buffer_turns": buffer.metadata.turn_count if buffer else 0,
|
||||
"buffer_tokens": buffer.metadata.total_tokens if buffer else 0
|
||||
}
|
||||
|
||||
# Get Qdrant stats
|
||||
qdrant_stats = await self.qdrant_memory.get_conversation_stats(conversation_id)
|
||||
|
||||
# Combine
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
**buffer_stats,
|
||||
"qdrant_turns": qdrant_stats["total_turns"],
|
||||
"qdrant_tokens": qdrant_stats["total_tokens"],
|
||||
"exists_in_buffer": buffer is not None,
|
||||
"exists_in_qdrant": qdrant_stats["exists"]
|
||||
}
|
||||
|
||||
|
||||
# Global instance
|
||||
_memory_manager: Optional[MemoryManager] = None
|
||||
|
||||
|
||||
def get_memory_manager() -> MemoryManager:
|
||||
"""
|
||||
Get or create global memory manager instance
|
||||
|
||||
Returns:
|
||||
MemoryManager instance
|
||||
"""
|
||||
global _memory_manager
|
||||
if _memory_manager is None:
|
||||
_memory_manager = MemoryManager()
|
||||
return _memory_manager
|
||||
@@ -1,389 +0,0 @@
|
||||
"""
|
||||
Unified Tier 2/3: Qdrant-based conversation memory
|
||||
|
||||
Single Qdrant collection serving both purposes:
|
||||
- Tier 2: Historical retrieval (filter by conversation_id, time-based)
|
||||
- Tier 3: Semantic search (vector similarity across conversations)
|
||||
"""
|
||||
import logging
|
||||
import uuid
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import (
|
||||
Distance,
|
||||
VectorParams,
|
||||
PointStruct,
|
||||
Filter,
|
||||
FieldCondition,
|
||||
MatchValue,
|
||||
Range,
|
||||
)
|
||||
|
||||
from .base import BaseMemory
|
||||
from .schemas import ConversationTurn, MessageRole
|
||||
from src.config import get_settings
|
||||
from src.models.embeddings_ollama import get_embedding_client
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class QdrantConversationMemory(BaseMemory):
|
||||
"""
|
||||
Unified conversation memory using Qdrant
|
||||
|
||||
Stores all conversation turns with vectors for semantic search.
|
||||
Can be queried in two ways:
|
||||
- Tier 2 mode: Filter by conversation_id for chronological history
|
||||
- Tier 3 mode: Vector similarity search for semantic recall
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
host: Optional[str] = None,
|
||||
port: Optional[int] = None
|
||||
):
|
||||
"""
|
||||
Initialize Qdrant memory
|
||||
|
||||
Args:
|
||||
collection_name: Name of Qdrant collection
|
||||
host: Qdrant host
|
||||
port: Qdrant port
|
||||
"""
|
||||
self.collection_name = collection_name or settings.qdrant_collection_conversations
|
||||
self.host = host or settings.qdrant_host
|
||||
self.port = port or settings.qdrant_port
|
||||
|
||||
# Initialize clients
|
||||
self.client = QdrantClient(host=self.host, port=self.port)
|
||||
self.embedding_client = get_embedding_client()
|
||||
|
||||
logger.info(
|
||||
f"Initialized QdrantConversationMemory: "
|
||||
f"{self.host}:{self.port}/{self.collection_name}"
|
||||
)
|
||||
|
||||
# Ensure collection exists
|
||||
self._ensure_collection()
|
||||
|
||||
def _ensure_collection(self) -> None:
|
||||
"""Create collection if it doesn't exist"""
|
||||
try:
|
||||
collections = self.client.get_collections().collections
|
||||
collection_names = [c.name for c in collections]
|
||||
|
||||
if self.collection_name not in collection_names:
|
||||
logger.info(f"Creating collection: {self.collection_name}")
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=VectorParams(
|
||||
size=settings.embedding_dimension,
|
||||
distance=Distance.COSINE
|
||||
)
|
||||
)
|
||||
logger.info(f"✓ Collection created: {self.collection_name}")
|
||||
else:
|
||||
logger.info(f"✓ Collection exists: {self.collection_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error ensuring collection: {e}")
|
||||
raise
|
||||
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a conversation turn with its embedding
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
# Generate embedding
|
||||
embedding = await self.embedding_client.embed_text(turn.content)
|
||||
|
||||
# Create point ID: deterministic UUID from conversation_id + turn_number
|
||||
# Qdrant requires UUID or unsigned int, so we generate UUID from string
|
||||
point_id_str = f"{conversation_id}_{turn.turn_number}"
|
||||
point_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, point_id_str))
|
||||
|
||||
# Build payload with user_id for multi-tenancy
|
||||
payload = {
|
||||
"conversation_id": conversation_id,
|
||||
"turn_number": turn.turn_number,
|
||||
"role": turn.role.value if isinstance(turn.role, MessageRole) else turn.role,
|
||||
"content": turn.content,
|
||||
"timestamp": turn.timestamp.isoformat(),
|
||||
"user_id": turn.user_id, # Multi-tenancy
|
||||
"metadata": turn.metadata,
|
||||
}
|
||||
|
||||
# Add token info if available
|
||||
if turn.tokens:
|
||||
payload["tokens_prompt"] = turn.tokens.prompt
|
||||
payload["tokens_completion"] = turn.tokens.completion
|
||||
payload["tokens_total"] = turn.tokens.total
|
||||
|
||||
# Upsert to Qdrant
|
||||
try:
|
||||
self.client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=[
|
||||
PointStruct(
|
||||
id=point_id,
|
||||
vector=embedding,
|
||||
payload=payload
|
||||
)
|
||||
]
|
||||
)
|
||||
logger.debug(f"Stored turn {turn.turn_number} for conversation {conversation_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing turn in Qdrant: {e}")
|
||||
raise
|
||||
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns for a conversation (Tier 2 mode: chronological)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
try:
|
||||
# Scroll through all points for this conversation
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=limit or 100,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
with_vectors=False
|
||||
)
|
||||
|
||||
# Convert to ConversationTurn objects
|
||||
turns = []
|
||||
for point in points:
|
||||
payload = point.payload
|
||||
turn = ConversationTurn(
|
||||
role=MessageRole(payload["role"]),
|
||||
content=payload["content"],
|
||||
timestamp=datetime.fromisoformat(payload["timestamp"]),
|
||||
turn_number=payload["turn_number"],
|
||||
user_id=payload.get("user_id", "llm-testuser"), # Multi-tenancy
|
||||
metadata=payload.get("metadata", {})
|
||||
)
|
||||
turns.append(turn)
|
||||
|
||||
# Sort by turn_number
|
||||
turns.sort(key=lambda t: t.turn_number)
|
||||
|
||||
return turns
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving turns from Qdrant: {e}")
|
||||
return []
|
||||
|
||||
async def similarity_search(
|
||||
self,
|
||||
query: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Semantic search for relevant turns (Tier 3 mode: semantic)
|
||||
|
||||
Args:
|
||||
query: Search query text
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
try:
|
||||
# Generate query embedding
|
||||
query_embedding = await self.embedding_client.embed_text(query)
|
||||
|
||||
# Build filter if conversation_id specified
|
||||
search_filter = None
|
||||
if conversation_id:
|
||||
search_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Search in Qdrant
|
||||
results = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_embedding,
|
||||
query_filter=search_filter,
|
||||
limit=limit,
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
# Convert results
|
||||
matches = []
|
||||
for result in results:
|
||||
payload = result.payload
|
||||
match = {
|
||||
"conversation_id": payload["conversation_id"],
|
||||
"turn_number": payload["turn_number"],
|
||||
"role": payload["role"],
|
||||
"content": payload["content"],
|
||||
"timestamp": payload["timestamp"],
|
||||
"score": result.score,
|
||||
}
|
||||
matches.append(match)
|
||||
|
||||
logger.debug(
|
||||
f"Semantic search found {len(matches)} matches for query: {query[:50]}..."
|
||||
)
|
||||
|
||||
return matches
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in semantic search: {e}")
|
||||
return []
|
||||
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
try:
|
||||
# Delete all points with this conversation_id
|
||||
self.client.delete(
|
||||
collection_name=self.collection_name,
|
||||
points_selector=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
logger.info(f"Cleared conversation {conversation_id} from Qdrant")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing conversation: {e}")
|
||||
raise
|
||||
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation has any turns
|
||||
"""
|
||||
try:
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=1,
|
||||
with_payload=False,
|
||||
with_vectors=False
|
||||
)
|
||||
return len(points) > 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking conversation existence: {e}")
|
||||
return False
|
||||
|
||||
async def get_conversation_stats(self, conversation_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get statistics about a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Dictionary with stats
|
||||
"""
|
||||
try:
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=1000, # Get all points
|
||||
with_payload=True,
|
||||
with_vectors=False
|
||||
)
|
||||
|
||||
total_turns = len(points)
|
||||
total_tokens = sum(
|
||||
point.payload.get("tokens_total", 0) for point in points
|
||||
)
|
||||
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
"total_turns": total_turns,
|
||||
"total_tokens": total_tokens,
|
||||
"exists": total_turns > 0
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting conversation stats: {e}")
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
"total_turns": 0,
|
||||
"total_tokens": 0,
|
||||
"exists": False
|
||||
}
|
||||
|
||||
|
||||
# Global instance
|
||||
_qdrant_memory: Optional[QdrantConversationMemory] = None
|
||||
|
||||
|
||||
def get_qdrant_memory() -> QdrantConversationMemory:
|
||||
"""
|
||||
Get or create global Qdrant memory instance
|
||||
|
||||
Returns:
|
||||
QdrantConversationMemory instance
|
||||
"""
|
||||
global _qdrant_memory
|
||||
if _qdrant_memory is None:
|
||||
_qdrant_memory = QdrantConversationMemory()
|
||||
return _qdrant_memory
|
||||
@@ -1,110 +0,0 @@
|
||||
"""
|
||||
Pydantic schemas for memory system
|
||||
"""
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MessageRole(str, Enum):
|
||||
"""Message role types"""
|
||||
SYSTEM = "system"
|
||||
USER = "user"
|
||||
ASSISTANT = "assistant"
|
||||
|
||||
|
||||
class TokenUsage(BaseModel):
|
||||
"""Token usage information"""
|
||||
prompt: int = 0
|
||||
completion: int = 0
|
||||
total: int = 0
|
||||
|
||||
|
||||
class ConversationTurn(BaseModel):
|
||||
"""A single turn in a conversation"""
|
||||
role: MessageRole
|
||||
content: str
|
||||
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
||||
turn_number: int
|
||||
user_id: str = "llm-testuser" # Multi-tenancy: user who owns this turn
|
||||
tokens: Optional[TokenUsage] = None
|
||||
metadata: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ConversationMetadata(BaseModel):
|
||||
"""Metadata about a conversation"""
|
||||
conversation_id: str
|
||||
user_id: str = "llm-testuser" # Multi-tenancy: user who owns this conversation
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
last_updated: datetime = Field(default_factory=datetime.utcnow)
|
||||
turn_count: int = 0
|
||||
total_tokens: int = 0
|
||||
status: str = "active" # active, archived, deleted
|
||||
|
||||
|
||||
class ConversationBuffer(BaseModel):
|
||||
"""In-memory conversation buffer (Tier 1)"""
|
||||
conversation_id: str
|
||||
turns: List[ConversationTurn] = Field(default_factory=list)
|
||||
metadata: ConversationMetadata
|
||||
|
||||
|
||||
class ConversationSummary(BaseModel):
|
||||
"""Summarized conversation segment (Tier 2)"""
|
||||
conversation_id: str
|
||||
summary_text: str
|
||||
turn_range_start: int
|
||||
turn_range_end: int
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
token_count: int = 0
|
||||
|
||||
|
||||
class MemoryQuery(BaseModel):
|
||||
"""Query for memory retrieval"""
|
||||
conversation_id: str
|
||||
query: Optional[str] = None
|
||||
limit: int = Field(default=10, ge=1, le=100)
|
||||
include_tier1: bool = True
|
||||
include_tier2: bool = True
|
||||
include_tier3: bool = True
|
||||
|
||||
|
||||
class MemoryResult(BaseModel):
|
||||
"""Result from memory retrieval"""
|
||||
conversation_id: str
|
||||
turns: List[ConversationTurn] = Field(default_factory=list)
|
||||
summaries: List[ConversationSummary] = Field(default_factory=list)
|
||||
source_tiers: List[int] = Field(default_factory=list) # Which tiers contributed
|
||||
total_results: int = 0
|
||||
|
||||
|
||||
# API Request/Response Models
|
||||
|
||||
class ConversationListResponse(BaseModel):
|
||||
"""Response for listing conversations"""
|
||||
conversations: List[ConversationMetadata]
|
||||
total: int
|
||||
page: int = 1
|
||||
page_size: int = 50
|
||||
|
||||
|
||||
class ConversationDetailResponse(BaseModel):
|
||||
"""Response for conversation details"""
|
||||
metadata: ConversationMetadata
|
||||
recent_turns: List[ConversationTurn]
|
||||
turn_count: int
|
||||
|
||||
|
||||
class ConversationSearchRequest(BaseModel):
|
||||
"""Request for semantic search in conversation"""
|
||||
query: str
|
||||
limit: int = Field(default=5, ge=1, le=50)
|
||||
|
||||
|
||||
class ConversationSearchResponse(BaseModel):
|
||||
"""Response for semantic search"""
|
||||
conversation_id: str
|
||||
results: List[ConversationTurn]
|
||||
scores: List[float] = Field(default_factory=list)
|
||||
total_results: int
|
||||
@@ -1,239 +0,0 @@
|
||||
"""
|
||||
Tier 1: ConversationBufferMemory (In-Memory Working Memory)
|
||||
|
||||
Fast in-memory storage for recent conversation turns.
|
||||
- Stores last N turns in RAM
|
||||
- < 1ms access time
|
||||
- Ephemeral (lost on restart)
|
||||
- Automatic pruning when limit reached
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
from datetime import datetime
|
||||
from collections import OrderedDict
|
||||
|
||||
from .base import Tier1Memory
|
||||
from .schemas import (
|
||||
ConversationTurn,
|
||||
ConversationBuffer,
|
||||
ConversationMetadata,
|
||||
MessageRole,
|
||||
TokenUsage
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConversationBufferMemory(Tier1Memory):
|
||||
"""
|
||||
In-memory buffer for recent conversation turns.
|
||||
|
||||
Stores the last N turns of each conversation in RAM for fast access.
|
||||
Automatically prunes old turns when limit is reached.
|
||||
"""
|
||||
|
||||
def __init__(self, max_turns: int = 10):
|
||||
"""
|
||||
Initialize buffer memory
|
||||
|
||||
Args:
|
||||
max_turns: Maximum number of turns to keep per conversation
|
||||
"""
|
||||
self.max_turns = max_turns
|
||||
# Use OrderedDict to maintain insertion order
|
||||
self._buffers: Dict[str, ConversationBuffer] = OrderedDict()
|
||||
logger.info(f"Initialized ConversationBufferMemory with max_turns={max_turns}")
|
||||
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a new turn to the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
# Get or create buffer
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
buffer = ConversationBuffer(
|
||||
conversation_id=conversation_id,
|
||||
turns=[],
|
||||
metadata=ConversationMetadata(
|
||||
conversation_id=conversation_id
|
||||
)
|
||||
)
|
||||
self._buffers[conversation_id] = buffer
|
||||
|
||||
# Add turn
|
||||
buffer.turns.append(turn)
|
||||
|
||||
# Update metadata
|
||||
buffer.metadata.turn_count = len(buffer.turns)
|
||||
buffer.metadata.last_updated = datetime.utcnow()
|
||||
|
||||
if turn.tokens:
|
||||
buffer.metadata.total_tokens += turn.tokens.total
|
||||
|
||||
# Auto-prune if exceeds max turns
|
||||
if len(buffer.turns) > self.max_turns:
|
||||
await self.prune(conversation_id, keep_last=self.max_turns)
|
||||
|
||||
logger.debug(
|
||||
f"Added turn {turn.turn_number} to conversation {conversation_id}. "
|
||||
f"Buffer size: {len(buffer.turns)}"
|
||||
)
|
||||
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns from the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return []
|
||||
|
||||
turns = buffer.turns[offset:]
|
||||
if limit:
|
||||
turns = turns[:limit]
|
||||
|
||||
return turns
|
||||
|
||||
async def get_recent_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 10
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get the most recent N turns
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Number of recent turns to retrieve
|
||||
|
||||
Returns:
|
||||
List of recent turns (most recent last)
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return []
|
||||
|
||||
return buffer.turns[-limit:] if len(buffer.turns) > limit else buffer.turns
|
||||
|
||||
async def get_buffer(self, conversation_id: str) -> Optional[ConversationBuffer]:
|
||||
"""
|
||||
Get the full conversation buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
ConversationBuffer or None if not found
|
||||
"""
|
||||
return self._buffers.get(conversation_id)
|
||||
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
if conversation_id in self._buffers:
|
||||
del self._buffers[conversation_id]
|
||||
logger.info(f"Cleared buffer for conversation {conversation_id}")
|
||||
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists in the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation exists
|
||||
"""
|
||||
return conversation_id in self._buffers
|
||||
|
||||
async def prune(self, conversation_id: str, keep_last: int = 5) -> None:
|
||||
"""
|
||||
Prune old turns, keeping only the most recent ones
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
keep_last: Number of recent turns to keep
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return
|
||||
|
||||
if len(buffer.turns) > keep_last:
|
||||
removed_count = len(buffer.turns) - keep_last
|
||||
buffer.turns = buffer.turns[-keep_last:]
|
||||
buffer.metadata.turn_count = len(buffer.turns)
|
||||
|
||||
logger.debug(
|
||||
f"Pruned {removed_count} turns from conversation {conversation_id}. "
|
||||
f"Kept last {keep_last} turns."
|
||||
)
|
||||
|
||||
async def get_all_conversation_ids(self) -> List[str]:
|
||||
"""
|
||||
Get list of all conversation IDs in memory
|
||||
|
||||
Returns:
|
||||
List of conversation IDs
|
||||
"""
|
||||
return list(self._buffers.keys())
|
||||
|
||||
async def get_buffer_stats(self) -> dict:
|
||||
"""
|
||||
Get statistics about buffer memory usage
|
||||
|
||||
Returns:
|
||||
Dictionary with stats
|
||||
"""
|
||||
total_conversations = len(self._buffers)
|
||||
total_turns = sum(len(buf.turns) for buf in self._buffers.values())
|
||||
total_tokens = sum(buf.metadata.total_tokens for buf in self._buffers.values())
|
||||
|
||||
return {
|
||||
"total_conversations": total_conversations,
|
||||
"total_turns": total_turns,
|
||||
"total_tokens": total_tokens,
|
||||
"max_turns_per_conversation": self.max_turns,
|
||||
"avg_turns_per_conversation": (
|
||||
total_turns / total_conversations if total_conversations > 0 else 0
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
# Global instance
|
||||
_buffer_memory: Optional[ConversationBufferMemory] = None
|
||||
|
||||
|
||||
def get_buffer_memory(max_turns: int = 10) -> ConversationBufferMemory:
|
||||
"""
|
||||
Get or create the global buffer memory instance
|
||||
|
||||
Args:
|
||||
max_turns: Maximum turns per conversation
|
||||
|
||||
Returns:
|
||||
ConversationBufferMemory instance
|
||||
"""
|
||||
global _buffer_memory
|
||||
if _buffer_memory is None:
|
||||
_buffer_memory = ConversationBufferMemory(max_turns=max_turns)
|
||||
return _buffer_memory
|
||||
Reference in New Issue
Block a user