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
|
||||
Reference in New Issue
Block a user