ai-flow improvement / add langchain
This commit is contained in:
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
Unified Agent Module
|
||||
|
||||
This module provides an intelligent agent that can handle infrastructure management,
|
||||
web search, and multi-step reasoning with transparent streaming output.
|
||||
"""
|
||||
from .orchestrator import UnifiedAgent, get_unified_agent
|
||||
from .tools import get_agent_tools, ALL_TOOLS
|
||||
|
||||
__all__ = [
|
||||
"UnifiedAgent",
|
||||
"get_unified_agent",
|
||||
"get_agent_tools",
|
||||
"ALL_TOOLS",
|
||||
]
|
||||
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
Agent Orchestrator - Unified intelligent agent with streaming reasoning
|
||||
|
||||
This orchestrator uses LangGraph to create a ReAct-style agent that can:
|
||||
- Use tools to answer infrastructure questions
|
||||
- Stream thinking/reasoning output
|
||||
- Handle multi-step tasks
|
||||
- Route to appropriate expert models
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
from typing import AsyncIterator, Dict, Any, List
|
||||
from functools import lru_cache
|
||||
|
||||
from langchain_ollama import ChatOllama
|
||||
from langgraph.prebuilt import create_react_agent
|
||||
from langgraph.graph import StateGraph
|
||||
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage
|
||||
|
||||
from src.config import get_settings
|
||||
from src.agent.tools import get_agent_tools
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class UnifiedAgent:
|
||||
"""
|
||||
Unified intelligent agent that handles all tool routing and reasoning
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.settings = get_settings()
|
||||
self.tools = get_agent_tools()
|
||||
|
||||
# Initialize Ollama LLM (must be a model that supports tool calling)
|
||||
self.llm = ChatOllama(
|
||||
model=self.settings.agent_model,
|
||||
base_url=self.settings.ollama_base_url,
|
||||
temperature=0.7,
|
||||
)
|
||||
|
||||
# Create ReAct agent with tools
|
||||
self.agent = create_react_agent(
|
||||
self.llm,
|
||||
self.tools,
|
||||
state_modifier=self._get_system_prompt(),
|
||||
)
|
||||
|
||||
logger.info(f"Initialized Unified Agent with {len(self.tools)} tools")
|
||||
|
||||
def _get_system_prompt(self) -> str:
|
||||
"""Get the system prompt that defines agent behavior"""
|
||||
return """You are Tatlock, a helpful personal assistant with the demeanor of a British butler.
|
||||
You address users as \"sir\" and speak formally.
|
||||
You are not overly apologetic and can be a little snarky at times.
|
||||
|
||||
Your capabilities:
|
||||
- Search the web and extract content
|
||||
- Monitor service health via Uptime Kuma
|
||||
- Read project documentation
|
||||
- Check system resources
|
||||
- Manage Docker containers and services via Portainer
|
||||
- Configure reverse proxies and domains via Nginx Proxy Manager
|
||||
|
||||
When helping users:
|
||||
1. Think step-by-step about what information you need
|
||||
2. Use tools when you need current/specific information
|
||||
3. Be concise but thorough in your responses
|
||||
4. If a task requires multiple steps, explain what you're doing
|
||||
5. Always verify information before making changes
|
||||
|
||||
Available infrastructure:
|
||||
- 22 running services (Ollama, Portainer, NPM, Jellyfin, Gitea, etc.)
|
||||
- GPU: NVIDIA RTX 2080 Ti (11GB VRAM)
|
||||
- Storage: SSD for configs, HDD for media
|
||||
- Network: Headscale mesh VPN + NPM reverse proxy
|
||||
|
||||
If you see an opportunity to make a pun or joke, you simply cannot resist.
|
||||
Be helpful, accurate, and transparent about what you're doing!"""
|
||||
|
||||
async def chat(
|
||||
self,
|
||||
message: str,
|
||||
conversation_history: List[Dict[str, str]] = None,
|
||||
stream: bool = True
|
||||
) -> 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
|
||||
|
||||
Yields:
|
||||
Dict with keys:
|
||||
- type: "thinking" | "tool_call" | "tool_result" | "content"
|
||||
- content: The actual content
|
||||
- tool: Tool name (if type is tool_call)
|
||||
- model: Model being used (optional)
|
||||
"""
|
||||
try:
|
||||
# Build message list
|
||||
messages = []
|
||||
|
||||
# Add conversation history if provided
|
||||
if conversation_history:
|
||||
for turn in conversation_history:
|
||||
if turn.get("role") == "user":
|
||||
messages.append(HumanMessage(content=turn["content"]))
|
||||
elif turn.get("role") == "assistant":
|
||||
messages.append(AIMessage(content=turn["content"]))
|
||||
|
||||
# Add current message
|
||||
messages.append(HumanMessage(content=message))
|
||||
|
||||
# Initial thinking
|
||||
yield {
|
||||
"type": "thinking",
|
||||
"content": "Analyzing your request...",
|
||||
"model": self.settings.default_model
|
||||
}
|
||||
|
||||
# Stream agent execution
|
||||
async for chunk in self.agent.astream(
|
||||
{"messages": messages},
|
||||
stream_mode="values" # Stream full state updates
|
||||
):
|
||||
# Extract messages from the chunk
|
||||
if "messages" in chunk:
|
||||
latest_messages = chunk["messages"]
|
||||
|
||||
# Process the latest message
|
||||
if latest_messages:
|
||||
latest = latest_messages[-1]
|
||||
|
||||
# Tool invocation
|
||||
if hasattr(latest, 'additional_kwargs') and 'tool_calls' in latest.additional_kwargs:
|
||||
tool_calls = latest.additional_kwargs['tool_calls']
|
||||
for tool_call in tool_calls:
|
||||
tool_name = tool_call.get('function', {}).get('name', 'unknown')
|
||||
yield {
|
||||
"type": "tool_call",
|
||||
"tool": tool_name,
|
||||
"content": f"Using tool: {tool_name}..."
|
||||
}
|
||||
|
||||
# Tool result
|
||||
elif isinstance(latest, ToolMessage):
|
||||
yield {
|
||||
"type": "tool_result",
|
||||
"content": "Tool execution complete"
|
||||
}
|
||||
|
||||
# AI response (final or intermediate)
|
||||
elif isinstance(latest, AIMessage) and latest.content:
|
||||
# Check if this is intermediate thinking or final response
|
||||
if hasattr(latest, 'additional_kwargs') and latest.additional_kwargs.get('tool_calls'):
|
||||
# This is thinking before a tool call
|
||||
yield {
|
||||
"type": "thinking",
|
||||
"content": latest.content
|
||||
}
|
||||
else:
|
||||
# This is the final response
|
||||
yield {
|
||||
"type": "content",
|
||||
"content": latest.content
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in 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
|
||||
) -> str:
|
||||
"""
|
||||
Get a non-streaming response (for backwards compatibility)
|
||||
|
||||
Args:
|
||||
message: User's message
|
||||
conversation_history: Previous conversation turns (optional)
|
||||
|
||||
Returns:
|
||||
The final response content
|
||||
"""
|
||||
final_content = ""
|
||||
async for chunk in self.chat(message, conversation_history, stream=True):
|
||||
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()
|
||||
@@ -0,0 +1,146 @@
|
||||
"""
|
||||
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
|
||||
tool_name = chunk.get("tool", "unknown")
|
||||
sse_chunk = {
|
||||
"id": request_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": "assistant",
|
||||
"content": f"[🔧 Using {tool_name}...]\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)
|
||||
# Split into words for smooth streaming
|
||||
words = content.split()
|
||||
for word in words:
|
||||
sse_chunk = {
|
||||
"id": request_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"content": word + " "
|
||||
},
|
||||
"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"
|
||||
@@ -0,0 +1,282 @@
|
||||
"""
|
||||
Agent Tools - LangChain-compatible tools for the unified agent
|
||||
|
||||
These tools wrap existing Core API functionality for use with LangGraph.
|
||||
"""
|
||||
from langchain_core.tools import tool
|
||||
from typing import List, Dict, Optional
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Infrastructure Management Tools
|
||||
# ============================================================================
|
||||
|
||||
@tool
|
||||
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)}"
|
||||
|
||||
|
||||
@tool
|
||||
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
|
||||
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
|
||||
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
|
||||
# ============================================================================
|
||||
|
||||
@tool
|
||||
async def web_search(url: str) -> str:
|
||||
"""
|
||||
Fetch and extract the main content from a web page.
|
||||
|
||||
Uses intelligent content extraction to get the most relevant text from articles,
|
||||
documentation, and blog posts. Perfect for answering questions that require current information.
|
||||
|
||||
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
|
||||
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
|
||||
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)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tool Registry
|
||||
# ============================================================================
|
||||
|
||||
# All available tools for the agent
|
||||
ALL_TOOLS = [
|
||||
list_services,
|
||||
get_service_details,
|
||||
list_domains,
|
||||
check_service_health,
|
||||
web_search,
|
||||
read_documentation,
|
||||
get_system_status,
|
||||
]
|
||||
|
||||
|
||||
def get_agent_tools() -> List:
|
||||
"""Get all tools available to the agent"""
|
||||
return ALL_TOOLS
|
||||
@@ -52,6 +52,7 @@ class Settings(BaseSettings):
|
||||
|
||||
# Model Configuration
|
||||
default_model: str = "gemma:7b"
|
||||
agent_model: str = "mistral:7b" # Must support tool calling
|
||||
lightweight_models: str = "gemma:2b,gemma:7b"
|
||||
heavy_models: str = "mistral:7b,gemma2:9b,mixtral:8x7b"
|
||||
code_models: str = "codestral:latest,codegemma:latest"
|
||||
@@ -73,9 +74,9 @@ class Settings(BaseSettings):
|
||||
qdrant_collection_documents: str = "core_api_documents"
|
||||
qdrant_collection_user_facts: str = "core_api_user_facts"
|
||||
|
||||
# Embeddings
|
||||
embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
|
||||
embedding_dimension: int = 384
|
||||
# Embeddings (using Ollama - no local models needed)
|
||||
embedding_model: str = "nomic-embed-text" # Ollama embedding model
|
||||
embedding_dimension: int = 768 # nomic-embed-text dimension
|
||||
embedding_batch_size: int = 32
|
||||
|
||||
# Infrastructure Management (from credentials.py)
|
||||
|
||||
@@ -31,6 +31,16 @@ 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__)
|
||||
|
||||
|
||||
@@ -294,6 +304,76 @@ class AIController(BaseController):
|
||||
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}")
|
||||
|
||||
# Extract conversation history
|
||||
history = []
|
||||
for msg in request.messages[:-1]: # All except last
|
||||
history.append({"role": msg.role.value, "content": msg.content})
|
||||
|
||||
# Get last message
|
||||
user_message = request.messages[-1].content
|
||||
|
||||
# Get agent
|
||||
agent = get_unified_agent()
|
||||
|
||||
# Stream response
|
||||
if request.stream:
|
||||
async def agent_stream_generator():
|
||||
agent_stream = agent.chat(
|
||||
message=user_message,
|
||||
conversation_history=history,
|
||||
stream=True
|
||||
)
|
||||
# Always use "Tatlock" as model name in responses
|
||||
async for sse_chunk in stream_agent_to_sse(agent_stream, request_id, "Tatlock"):
|
||||
yield sse_chunk
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
# 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:
|
||||
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:
|
||||
@@ -385,25 +465,17 @@ class AIController(BaseController):
|
||||
)
|
||||
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)
|
||||
# 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(
|
||||
|
||||
@@ -730,6 +730,44 @@ class InfrastructureController(BaseController):
|
||||
logger.error(f"Failed to create proxy host: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@router.put(
|
||||
"/proxy/{proxy_id}",
|
||||
response_model=OperationResult,
|
||||
summary="Update a proxy host",
|
||||
description="Update an existing Nginx Proxy Manager proxy host configuration. Requires admin authentication."
|
||||
)
|
||||
async def update_proxy(
|
||||
proxy_id: int,
|
||||
config: Dict[str, Any],
|
||||
user: Dict = Depends(get_admin_user)
|
||||
):
|
||||
"""
|
||||
Update an existing Nginx Proxy Manager proxy host
|
||||
|
||||
Args:
|
||||
proxy_id: Proxy host ID to update
|
||||
config: Full proxy host configuration (get from get_proxy_host, modify, then update)
|
||||
|
||||
Returns:
|
||||
Operation result with updated proxy host details
|
||||
"""
|
||||
npm = get_npm_client()
|
||||
|
||||
try:
|
||||
result = await npm.update_proxy_host(proxy_id, config)
|
||||
|
||||
logger.info(f"Updated proxy host {proxy_id}: {result.get('domain_names', [])}")
|
||||
|
||||
return OperationResult(
|
||||
success=True,
|
||||
message=f"Proxy host {proxy_id} updated successfully",
|
||||
details={"proxy_host": result}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update proxy host {proxy_id}: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# Service Control Endpoints
|
||||
@router.get(
|
||||
"/service-groups",
|
||||
|
||||
@@ -23,7 +23,7 @@ from qdrant_client.models import (
|
||||
from .base import BaseMemory
|
||||
from .schemas import ConversationTurn, MessageRole
|
||||
from src.config import get_settings
|
||||
from src.models.embeddings import get_embedding_client
|
||||
from src.models.embeddings_ollama import get_embedding_client
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
@@ -101,7 +101,7 @@ class QdrantConversationMemory(BaseMemory):
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
# Generate embedding
|
||||
embedding = self.embedding_client.embed_text(turn.content)
|
||||
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
|
||||
@@ -218,7 +218,7 @@ class QdrantConversationMemory(BaseMemory):
|
||||
"""
|
||||
try:
|
||||
# Generate query embedding
|
||||
query_embedding = self.embedding_client.embed_text(query)
|
||||
query_embedding = await self.embedding_client.embed_text(query)
|
||||
|
||||
# Build filter if conversation_id specified
|
||||
search_filter = None
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
"""
|
||||
Ollama-based embedding client for text vectorization
|
||||
|
||||
Uses Ollama's embedding API instead of local sentence-transformers.
|
||||
This eliminates the need for PyTorch and heavy ML dependencies.
|
||||
"""
|
||||
import logging
|
||||
import httpx
|
||||
from typing import List, Optional
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class OllamaEmbeddingClient:
|
||||
"""Client for generating text embeddings using Ollama"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
timeout: int = 30
|
||||
):
|
||||
"""
|
||||
Initialize Ollama embedding client
|
||||
|
||||
Args:
|
||||
model_name: Embedding model name (default: nomic-embed-text)
|
||||
base_url: Ollama base URL (default from settings)
|
||||
timeout: Request timeout in seconds
|
||||
"""
|
||||
self.model_name = model_name or settings.embedding_model
|
||||
self.base_url = (base_url or settings.ollama_base_url).rstrip("/")
|
||||
self.timeout = timeout
|
||||
self.dimension = settings.embedding_dimension
|
||||
|
||||
logger.info(f"Initializing OllamaEmbeddingClient with model: {self.model_name}")
|
||||
logger.info(f"Ollama URL: {self.base_url}")
|
||||
|
||||
async def embed_text(self, text: str) -> List[float]:
|
||||
"""
|
||||
Generate embedding for a single text using Ollama
|
||||
|
||||
Args:
|
||||
text: Input text to embed
|
||||
|
||||
Returns:
|
||||
List of floats representing the embedding vector
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/embeddings",
|
||||
json={
|
||||
"model": self.model_name,
|
||||
"prompt": text
|
||||
}
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
return result["embedding"]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating embedding via Ollama: {e}")
|
||||
raise
|
||||
|
||||
async def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts
|
||||
|
||||
Args:
|
||||
texts: List of input texts
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
embeddings = []
|
||||
for text in texts:
|
||||
embedding = await self.embed_text(text)
|
||||
embeddings.append(embedding)
|
||||
return embeddings
|
||||
|
||||
def get_dimension(self) -> int:
|
||||
"""
|
||||
Get embedding dimension
|
||||
|
||||
Returns:
|
||||
Embedding vector dimension
|
||||
"""
|
||||
return self.dimension
|
||||
|
||||
|
||||
# Global instance
|
||||
_embedding_client: Optional[OllamaEmbeddingClient] = None
|
||||
|
||||
|
||||
def get_embedding_client() -> OllamaEmbeddingClient:
|
||||
"""
|
||||
Get or create global Ollama embedding client instance
|
||||
|
||||
Returns:
|
||||
OllamaEmbeddingClient instance
|
||||
"""
|
||||
global _embedding_client
|
||||
if _embedding_client is None:
|
||||
_embedding_client = OllamaEmbeddingClient()
|
||||
return _embedding_client
|
||||
|
||||
|
||||
async def embed_text_async(text: str) -> List[float]:
|
||||
"""
|
||||
Async wrapper for embedding text
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
|
||||
Returns:
|
||||
Embedding vector
|
||||
"""
|
||||
client = get_embedding_client()
|
||||
return await client.embed_text(text)
|
||||
|
||||
|
||||
async def embed_batch_async(texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Async wrapper for batch embedding
|
||||
|
||||
Args:
|
||||
texts: List of input texts
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
client = get_embedding_client()
|
||||
return await client.embed_batch(texts)
|
||||
Reference in New Issue
Block a user