ai-flow improvement / add langchain

This commit is contained in:
2025-11-23 14:51:19 +01:00
parent ade84f34d5
commit 5e734ad27f
25 changed files with 4867 additions and 51 deletions
+15
View File
@@ -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",
]
+204
View File
@@ -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()
+146
View File
@@ -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"
+282
View File
@@ -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