5290287058
- TTS: xtts-v2 integration with voice cloning support
- Document: docling integration for PDF/DOCX/PPTX processing
- Memory Service: added /facts/upsert, /facts/{key}, /facts endpoints
- Added required dependencies (TTS, docling)
924 lines
33 KiB
Python
924 lines
33 KiB
Python
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
from typing import Literal, Optional, Dict, Any, List
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import yaml
|
|
import httpx
|
|
import logging
|
|
from neo4j import AsyncGraphDatabase
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
app = FastAPI(title="DAARION Router", version="2.0.0")
|
|
|
|
# Configuration
|
|
NATS_URL = os.getenv("NATS_URL", "nats://nats:4222")
|
|
SWAPPER_URL = os.getenv("SWAPPER_URL", "http://swapper-service:8890")
|
|
# All multimodal services now through Swapper
|
|
STT_URL = os.getenv("STT_URL", "http://swapper-service:8890") # Swapper /stt endpoint
|
|
TTS_URL = os.getenv("TTS_URL", "http://swapper-service:8890") # Swapper /tts endpoint
|
|
VISION_URL = os.getenv("VISION_URL", "http://172.18.0.1:11434") # Host Ollama
|
|
OCR_URL = os.getenv("OCR_URL", "http://swapper-service:8890") # Swapper /ocr endpoint
|
|
DOCUMENT_URL = os.getenv("DOCUMENT_URL", "http://swapper-service:8890") # Swapper /document endpoint
|
|
CITY_SERVICE_URL = os.getenv("CITY_SERVICE_URL", "http://daarion-city-service:7001")
|
|
|
|
# Neo4j Configuration
|
|
NEO4J_URI = os.getenv("NEO4J_BOLT_URL", "bolt://neo4j:7687")
|
|
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
|
|
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "DaarionNeo4j2026!")
|
|
|
|
# HTTP client for backend services
|
|
http_client: Optional[httpx.AsyncClient] = None
|
|
|
|
# Neo4j driver
|
|
neo4j_driver = None
|
|
neo4j_available = False
|
|
|
|
# NATS client
|
|
nc = None
|
|
nats_available = False
|
|
|
|
# Models
|
|
class FilterDecision(BaseModel):
|
|
channel_id: str
|
|
message_id: Optional[str] = None
|
|
matrix_event_id: str
|
|
microdao_id: str
|
|
decision: Literal["allow", "deny", "modify"]
|
|
target_agent_id: Optional[str] = None
|
|
rewrite_prompt: Optional[str] = None
|
|
|
|
class AgentInvocation(BaseModel):
|
|
agent_id: str
|
|
entrypoint: Literal["channel_message", "direct", "cron"] = "channel_message"
|
|
payload: Dict[str, Any]
|
|
|
|
# Load config
|
|
def load_config():
|
|
config_path = "router_config.yaml"
|
|
if os.path.exists(config_path):
|
|
with open(config_path, 'r') as f:
|
|
return yaml.safe_load(f)
|
|
return {
|
|
"messaging_inbound": {
|
|
"enabled": True,
|
|
"source_subject": "agent.filter.decision",
|
|
"target_subject": "router.invoke.agent"
|
|
}
|
|
}
|
|
|
|
def load_router_config():
|
|
"""Load main router-config.yml with agents and LLM profiles"""
|
|
# Try multiple locations
|
|
paths = [
|
|
"router-config.yml",
|
|
"/app/router-config.yml",
|
|
"../router-config.yml",
|
|
"../../router-config.yml"
|
|
]
|
|
|
|
for path in paths:
|
|
if os.path.exists(path):
|
|
with open(path, 'r') as f:
|
|
logger.info(f"✅ Loaded router config from {path}")
|
|
return yaml.safe_load(f)
|
|
|
|
logger.warning("⚠️ router-config.yml not found, using empty config")
|
|
return {"agents": {}}
|
|
|
|
config = load_config()
|
|
router_config = load_router_config()
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
"""Initialize NATS connection and subscriptions"""
|
|
global nc, nats_available, http_client, neo4j_driver, neo4j_available
|
|
logger.info("🚀 DAGI Router v2.0.0 starting up...")
|
|
|
|
# Initialize HTTP client
|
|
http_client = httpx.AsyncClient(timeout=60.0)
|
|
logger.info("✅ HTTP client initialized")
|
|
|
|
# Initialize Neo4j connection
|
|
try:
|
|
neo4j_driver = AsyncGraphDatabase.driver(
|
|
NEO4J_URI,
|
|
auth=(NEO4J_USER, NEO4J_PASSWORD)
|
|
)
|
|
# Verify connection
|
|
async with neo4j_driver.session() as session:
|
|
result = await session.run("RETURN 1 as test")
|
|
await result.consume()
|
|
neo4j_available = True
|
|
logger.info(f"✅ Connected to Neo4j at {NEO4J_URI}")
|
|
except Exception as e:
|
|
logger.warning(f"⚠️ Neo4j not available: {e}")
|
|
neo4j_available = False
|
|
|
|
# Try to connect to NATS
|
|
try:
|
|
import nats
|
|
nc = await nats.connect(NATS_URL)
|
|
nats_available = True
|
|
logger.info(f"✅ Connected to NATS at {NATS_URL}")
|
|
|
|
# Subscribe to filter decisions if enabled
|
|
if config.get("messaging_inbound", {}).get("enabled", True):
|
|
asyncio.create_task(subscribe_to_filter_decisions())
|
|
else:
|
|
logger.warning("⚠️ Messaging inbound routing disabled in config")
|
|
except Exception as e:
|
|
logger.warning(f"⚠️ NATS not available: {e}")
|
|
logger.warning("⚠️ Running in test mode (HTTP only)")
|
|
nats_available = False
|
|
|
|
# Log backend URLs
|
|
logger.info(f"📡 Swapper URL: {SWAPPER_URL}")
|
|
logger.info(f"📡 STT URL: {STT_URL}")
|
|
logger.info(f"📡 Vision URL: {VISION_URL}")
|
|
logger.info(f"📡 OCR URL: {OCR_URL}")
|
|
logger.info(f"📡 Neo4j URL: {NEO4J_URI}")
|
|
|
|
async def subscribe_to_filter_decisions():
|
|
"""Subscribe to agent.filter.decision events"""
|
|
if not nc:
|
|
return
|
|
|
|
source_subject = config.get("messaging_inbound", {}).get(
|
|
"source_subject",
|
|
"agent.filter.decision"
|
|
)
|
|
|
|
try:
|
|
sub = await nc.subscribe(source_subject)
|
|
print(f"✅ Subscribed to {source_subject}")
|
|
|
|
async for msg in sub.messages:
|
|
try:
|
|
decision_data = json.loads(msg.data.decode())
|
|
await handle_filter_decision(decision_data)
|
|
except Exception as e:
|
|
print(f"❌ Error processing decision: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
except Exception as e:
|
|
print(f"❌ Subscription error: {e}")
|
|
|
|
async def handle_filter_decision(decision_data: dict):
|
|
"""
|
|
Process agent.filter.decision events
|
|
|
|
Only processes 'allow' decisions
|
|
Creates AgentInvocation and publishes to router.invoke.agent
|
|
"""
|
|
try:
|
|
print(f"\n🔀 Processing filter decision")
|
|
decision = FilterDecision(**decision_data)
|
|
|
|
# Only process 'allow' decisions
|
|
if decision.decision != "allow":
|
|
print(f"⏭️ Ignoring non-allow decision: {decision.decision}")
|
|
return
|
|
|
|
if not decision.target_agent_id:
|
|
print(f"⚠️ No target agent specified, skipping")
|
|
return
|
|
|
|
print(f"✅ Decision: allow")
|
|
print(f"📝 Target: {decision.target_agent_id}")
|
|
print(f"📝 Channel: {decision.channel_id}")
|
|
|
|
# Create AgentInvocation
|
|
invocation = AgentInvocation(
|
|
agent_id=decision.target_agent_id,
|
|
entrypoint="channel_message",
|
|
payload={
|
|
"channel_id": decision.channel_id,
|
|
"message_id": decision.message_id,
|
|
"matrix_event_id": decision.matrix_event_id,
|
|
"microdao_id": decision.microdao_id,
|
|
"rewrite_prompt": decision.rewrite_prompt
|
|
}
|
|
)
|
|
|
|
print(f"🚀 Created invocation for {invocation.agent_id}")
|
|
|
|
# Publish to NATS
|
|
await publish_agent_invocation(invocation)
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error handling decision: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
async def publish_agent_invocation(invocation: AgentInvocation):
|
|
"""Publish AgentInvocation to router.invoke.agent"""
|
|
if nc and nats_available:
|
|
target_subject = config.get("messaging_inbound", {}).get(
|
|
"target_subject",
|
|
"router.invoke.agent"
|
|
)
|
|
|
|
try:
|
|
await nc.publish(target_subject, invocation.json().encode())
|
|
print(f"✅ Published invocation to {target_subject}")
|
|
except Exception as e:
|
|
print(f"❌ Error publishing to NATS: {e}")
|
|
else:
|
|
print(f"⚠️ NATS not available, invocation not published: {invocation.json()}")
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
"""Health check endpoint"""
|
|
return {
|
|
"status": "ok",
|
|
"service": "router",
|
|
"version": "1.0.0",
|
|
"nats_connected": nats_available,
|
|
"messaging_inbound_enabled": config.get("messaging_inbound", {}).get("enabled", True)
|
|
}
|
|
|
|
@app.post("/internal/router/test-messaging", response_model=AgentInvocation)
|
|
async def test_messaging_route(decision: FilterDecision):
|
|
"""
|
|
Test endpoint for routing logic
|
|
|
|
Tests filter decision → agent invocation mapping without NATS
|
|
"""
|
|
print(f"\n🧪 Test routing request")
|
|
|
|
if decision.decision != "allow" or not decision.target_agent_id:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Decision not routable: {decision.decision}, agent: {decision.target_agent_id}"
|
|
)
|
|
|
|
invocation = AgentInvocation(
|
|
agent_id=decision.target_agent_id,
|
|
entrypoint="channel_message",
|
|
payload={
|
|
"channel_id": decision.channel_id,
|
|
"message_id": decision.message_id,
|
|
"matrix_event_id": decision.matrix_event_id,
|
|
"microdao_id": decision.microdao_id,
|
|
"rewrite_prompt": decision.rewrite_prompt
|
|
}
|
|
)
|
|
|
|
print(f"✅ Test invocation created for {invocation.agent_id}")
|
|
return invocation
|
|
|
|
@app.on_event("shutdown")
|
|
async def shutdown_event():
|
|
"""Clean shutdown"""
|
|
global nc, http_client
|
|
if nc:
|
|
await nc.close()
|
|
logger.info("✅ NATS connection closed")
|
|
if http_client:
|
|
await http_client.aclose()
|
|
logger.info("✅ HTTP client closed")
|
|
|
|
|
|
# ============================================================================
|
|
# Backend Integration Endpoints
|
|
# ============================================================================
|
|
|
|
class InferRequest(BaseModel):
|
|
"""Request for agent inference"""
|
|
prompt: str
|
|
model: Optional[str] = None
|
|
max_tokens: Optional[int] = 2048
|
|
temperature: Optional[float] = 0.7
|
|
system_prompt: Optional[str] = None
|
|
|
|
|
|
class InferResponse(BaseModel):
|
|
"""Response from agent inference"""
|
|
response: str
|
|
model: str
|
|
tokens_used: Optional[int] = None
|
|
backend: str
|
|
|
|
|
|
class BackendStatus(BaseModel):
|
|
"""Status of a backend service"""
|
|
name: str
|
|
url: str
|
|
status: str # online, offline, error
|
|
active_model: Optional[str] = None
|
|
error: Optional[str] = None
|
|
|
|
|
|
@app.get("/backends/status", response_model=List[BackendStatus])
|
|
async def get_backends_status():
|
|
"""Get status of all backend services"""
|
|
backends = []
|
|
|
|
# Check Swapper
|
|
try:
|
|
resp = await http_client.get(f"{SWAPPER_URL}/health", timeout=5.0)
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
backends.append(BackendStatus(
|
|
name="swapper",
|
|
url=SWAPPER_URL,
|
|
status="online",
|
|
active_model=data.get("active_model")
|
|
))
|
|
else:
|
|
backends.append(BackendStatus(
|
|
name="swapper",
|
|
url=SWAPPER_URL,
|
|
status="error",
|
|
error=f"HTTP {resp.status_code}"
|
|
))
|
|
except Exception as e:
|
|
backends.append(BackendStatus(
|
|
name="swapper",
|
|
url=SWAPPER_URL,
|
|
status="offline",
|
|
error=str(e)
|
|
))
|
|
|
|
# Check STT
|
|
try:
|
|
resp = await http_client.get(f"{STT_URL}/health", timeout=5.0)
|
|
backends.append(BackendStatus(
|
|
name="stt",
|
|
url=STT_URL,
|
|
status="online" if resp.status_code == 200 else "error"
|
|
))
|
|
except Exception as e:
|
|
backends.append(BackendStatus(
|
|
name="stt",
|
|
url=STT_URL,
|
|
status="offline",
|
|
error=str(e)
|
|
))
|
|
|
|
# Check Vision (Ollama)
|
|
try:
|
|
resp = await http_client.get(f"{VISION_URL}/api/tags", timeout=5.0)
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
models = [m.get("name") for m in data.get("models", [])]
|
|
backends.append(BackendStatus(
|
|
name="vision",
|
|
url=VISION_URL,
|
|
status="online",
|
|
active_model=", ".join(models[:3]) if models else None
|
|
))
|
|
else:
|
|
backends.append(BackendStatus(
|
|
name="vision",
|
|
url=VISION_URL,
|
|
status="error"
|
|
))
|
|
except Exception as e:
|
|
backends.append(BackendStatus(
|
|
name="vision",
|
|
url=VISION_URL,
|
|
status="offline",
|
|
error=str(e)
|
|
))
|
|
|
|
# Check OCR
|
|
try:
|
|
resp = await http_client.get(f"{OCR_URL}/health", timeout=5.0)
|
|
backends.append(BackendStatus(
|
|
name="ocr",
|
|
url=OCR_URL,
|
|
status="online" if resp.status_code == 200 else "error"
|
|
))
|
|
except Exception as e:
|
|
backends.append(BackendStatus(
|
|
name="ocr",
|
|
url=OCR_URL,
|
|
status="offline",
|
|
error=str(e)
|
|
))
|
|
|
|
return backends
|
|
|
|
|
|
@app.post("/v1/agents/{agent_id}/infer", response_model=InferResponse)
|
|
async def agent_infer(agent_id: str, request: InferRequest):
|
|
"""
|
|
Route inference request to appropriate backend.
|
|
|
|
Router decides which backend to use based on:
|
|
- Agent configuration (model, capabilities)
|
|
- Request type (text, vision, audio)
|
|
- Backend availability
|
|
|
|
System prompt is fetched from database via city-service API.
|
|
"""
|
|
logger.info(f"🔀 Inference request for agent: {agent_id}")
|
|
logger.info(f"📝 Prompt: {request.prompt[:100]}...")
|
|
|
|
# Get system prompt from database or config
|
|
system_prompt = request.system_prompt
|
|
|
|
if not system_prompt:
|
|
try:
|
|
from prompt_builder import get_agent_system_prompt
|
|
system_prompt = await get_agent_system_prompt(
|
|
agent_id,
|
|
city_service_url=CITY_SERVICE_URL,
|
|
router_config=router_config
|
|
)
|
|
logger.info(f"✅ Loaded system prompt from database for {agent_id}")
|
|
except Exception as e:
|
|
logger.warning(f"⚠️ Could not load prompt from database: {e}")
|
|
# Fallback to config
|
|
agent_config = router_config.get("agents", {}).get(agent_id, {})
|
|
system_prompt = agent_config.get("system_prompt")
|
|
|
|
# Determine which backend to use
|
|
# Use router config to get default model for agent, fallback to qwen3-8b
|
|
agent_config = router_config.get("agents", {}).get(agent_id, {})
|
|
default_llm = agent_config.get("default_llm", "qwen3-8b")
|
|
|
|
# Check if there's a routing rule for this agent
|
|
routing_rules = router_config.get("routing", [])
|
|
for rule in routing_rules:
|
|
if rule.get("when", {}).get("agent") == agent_id:
|
|
if "use_llm" in rule:
|
|
default_llm = rule.get("use_llm")
|
|
logger.info(f"🎯 Agent {agent_id} routing to: {default_llm}")
|
|
break
|
|
|
|
# Get LLM profile configuration
|
|
llm_profiles = router_config.get("llm_profiles", {})
|
|
llm_profile = llm_profiles.get(default_llm, {})
|
|
provider = llm_profile.get("provider", "ollama")
|
|
|
|
# Determine model name
|
|
if provider in ["deepseek", "openai", "anthropic", "mistral"]:
|
|
model = llm_profile.get("model", "deepseek-chat")
|
|
else:
|
|
# For local ollama, use swapper model name format
|
|
model = request.model or "qwen3-8b"
|
|
|
|
# =========================================================================
|
|
# CLOUD PROVIDERS (DeepSeek, OpenAI, etc.)
|
|
# =========================================================================
|
|
if provider == "deepseek":
|
|
try:
|
|
api_key = os.getenv(llm_profile.get("api_key_env", "DEEPSEEK_API_KEY"))
|
|
base_url = llm_profile.get("base_url", "https://api.deepseek.com")
|
|
|
|
if not api_key:
|
|
logger.error("❌ DeepSeek API key not configured")
|
|
raise HTTPException(status_code=500, detail="DeepSeek API key not configured")
|
|
|
|
logger.info(f"🌐 Calling DeepSeek API with model: {model}")
|
|
|
|
# Build messages array for chat completion
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
messages.append({"role": "user", "content": request.prompt})
|
|
|
|
deepseek_resp = await http_client.post(
|
|
f"{base_url}/v1/chat/completions",
|
|
headers={
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json"
|
|
},
|
|
json={
|
|
"model": model,
|
|
"messages": messages,
|
|
"max_tokens": request.max_tokens or llm_profile.get("max_tokens", 2048),
|
|
"temperature": request.temperature or llm_profile.get("temperature", 0.2),
|
|
"stream": False
|
|
},
|
|
timeout=llm_profile.get("timeout_ms", 40000) / 1000
|
|
)
|
|
|
|
if deepseek_resp.status_code == 200:
|
|
data = deepseek_resp.json()
|
|
response_text = data.get("choices", [{}])[0].get("message", {}).get("content", "")
|
|
tokens_used = data.get("usage", {}).get("total_tokens", 0)
|
|
|
|
logger.info(f"✅ DeepSeek response received, {tokens_used} tokens")
|
|
return InferResponse(
|
|
response=response_text,
|
|
model=model,
|
|
tokens_used=tokens_used,
|
|
backend="deepseek-cloud"
|
|
)
|
|
else:
|
|
logger.error(f"❌ DeepSeek error: {deepseek_resp.status_code} - {deepseek_resp.text}")
|
|
raise HTTPException(status_code=deepseek_resp.status_code, detail=f"DeepSeek API error: {deepseek_resp.text}")
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"❌ DeepSeek error: {e}")
|
|
# Don't fallback to local for cloud agents - raise error
|
|
raise HTTPException(status_code=503, detail=f"DeepSeek API error: {str(e)}")
|
|
|
|
# =========================================================================
|
|
# LOCAL PROVIDERS (Ollama via Swapper)
|
|
# =========================================================================
|
|
try:
|
|
# Check if Swapper is available
|
|
health_resp = await http_client.get(f"{SWAPPER_URL}/health", timeout=5.0)
|
|
if health_resp.status_code == 200:
|
|
logger.info(f"📡 Calling Swapper with model: {model}")
|
|
# Generate response via Swapper (which handles model loading)
|
|
generate_resp = await http_client.post(
|
|
f"{SWAPPER_URL}/generate",
|
|
json={
|
|
"model": model,
|
|
"prompt": request.prompt,
|
|
"system": system_prompt,
|
|
"max_tokens": request.max_tokens,
|
|
"temperature": request.temperature,
|
|
"stream": False
|
|
},
|
|
timeout=300.0
|
|
)
|
|
|
|
if generate_resp.status_code == 200:
|
|
data = generate_resp.json()
|
|
return InferResponse(
|
|
response=data.get("response", ""),
|
|
model=model,
|
|
tokens_used=data.get("eval_count", 0),
|
|
backend="swapper+ollama"
|
|
)
|
|
else:
|
|
logger.error(f"❌ Swapper error: {generate_resp.status_code} - {generate_resp.text}")
|
|
except Exception as e:
|
|
logger.error(f"❌ Swapper/Ollama error: {e}")
|
|
# Fallback to direct Ollama if Swapper fails
|
|
try:
|
|
logger.info(f"🔄 Falling back to direct Ollama connection")
|
|
generate_resp = await http_client.post(
|
|
f"{VISION_URL}/api/generate",
|
|
json={
|
|
"model": "qwen3:8b", # Use actual Ollama model name
|
|
"prompt": request.prompt,
|
|
"system": system_prompt,
|
|
"stream": False,
|
|
"options": {
|
|
"num_predict": request.max_tokens,
|
|
"temperature": request.temperature
|
|
}
|
|
},
|
|
timeout=120.0
|
|
)
|
|
|
|
if generate_resp.status_code == 200:
|
|
data = generate_resp.json()
|
|
return InferResponse(
|
|
response=data.get("response", ""),
|
|
model=model,
|
|
tokens_used=data.get("eval_count", 0),
|
|
backend="ollama-direct"
|
|
)
|
|
except Exception as e2:
|
|
logger.error(f"❌ Direct Ollama fallback also failed: {e2}")
|
|
|
|
# Fallback: return error
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail=f"No backend available for model: {model}"
|
|
)
|
|
|
|
|
|
@app.get("/v1/models")
|
|
async def list_available_models():
|
|
"""List all available models across backends"""
|
|
models = []
|
|
|
|
# Get Swapper models
|
|
try:
|
|
resp = await http_client.get(f"{SWAPPER_URL}/models", timeout=5.0)
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
for m in data.get("models", []):
|
|
models.append({
|
|
"id": m.get("name"),
|
|
"backend": "swapper",
|
|
"size_gb": m.get("size_gb"),
|
|
"status": m.get("status", "available")
|
|
})
|
|
except Exception as e:
|
|
logger.warning(f"Cannot get Swapper models: {e}")
|
|
|
|
# Get Ollama models
|
|
try:
|
|
resp = await http_client.get(f"{VISION_URL}/api/tags", timeout=5.0)
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
for m in data.get("models", []):
|
|
# Avoid duplicates
|
|
model_name = m.get("name")
|
|
if not any(x.get("id") == model_name for x in models):
|
|
models.append({
|
|
"id": model_name,
|
|
"backend": "ollama",
|
|
"size_gb": round(m.get("size", 0) / 1e9, 1),
|
|
"status": "loaded"
|
|
})
|
|
except Exception as e:
|
|
logger.warning(f"Cannot get Ollama models: {e}")
|
|
|
|
return {"models": models, "total": len(models)}
|
|
|
|
|
|
# =============================================================================
|
|
# NEO4J GRAPH API ENDPOINTS
|
|
# =============================================================================
|
|
|
|
class GraphNode(BaseModel):
|
|
"""Model for creating/updating a graph node"""
|
|
label: str # Node type: User, Agent, Topic, Fact, Entity, etc.
|
|
properties: Dict[str, Any]
|
|
node_id: Optional[str] = None # If provided, update existing node
|
|
|
|
class GraphRelationship(BaseModel):
|
|
"""Model for creating a relationship between nodes"""
|
|
from_node_id: str
|
|
to_node_id: str
|
|
relationship_type: str # KNOWS, MENTIONED, RELATED_TO, CREATED_BY, etc.
|
|
properties: Optional[Dict[str, Any]] = None
|
|
|
|
class GraphQuery(BaseModel):
|
|
"""Model for querying the graph"""
|
|
cypher: Optional[str] = None # Direct Cypher query (advanced)
|
|
# Or use structured query:
|
|
node_label: Optional[str] = None
|
|
node_id: Optional[str] = None
|
|
relationship_type: Optional[str] = None
|
|
depth: int = 1 # How many hops to traverse
|
|
limit: int = 50
|
|
|
|
class GraphSearchRequest(BaseModel):
|
|
"""Natural language search in graph"""
|
|
query: str
|
|
entity_types: Optional[List[str]] = None # Filter by types
|
|
limit: int = 20
|
|
|
|
|
|
@app.post("/v1/graph/nodes")
|
|
async def create_graph_node(node: GraphNode):
|
|
"""Create or update a node in the knowledge graph"""
|
|
if not neo4j_available or not neo4j_driver:
|
|
raise HTTPException(status_code=503, detail="Neo4j not available")
|
|
|
|
try:
|
|
async with neo4j_driver.session() as session:
|
|
# Generate node_id if not provided
|
|
node_id = node.node_id or f"{node.label.lower()}_{os.urandom(8).hex()}"
|
|
|
|
# Build properties with node_id
|
|
props = {**node.properties, "node_id": node_id, "updated_at": "datetime()"}
|
|
|
|
# Create or merge node
|
|
cypher = f"""
|
|
MERGE (n:{node.label} {{node_id: $node_id}})
|
|
SET n += $properties
|
|
SET n.updated_at = datetime()
|
|
RETURN n
|
|
"""
|
|
|
|
result = await session.run(cypher, node_id=node_id, properties=node.properties)
|
|
record = await result.single()
|
|
|
|
if record:
|
|
created_node = dict(record["n"])
|
|
logger.info(f"📊 Created/updated node: {node.label} - {node_id}")
|
|
return {"status": "ok", "node_id": node_id, "node": created_node}
|
|
|
|
raise HTTPException(status_code=500, detail="Failed to create node")
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Neo4j error creating node: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/v1/graph/relationships")
|
|
async def create_graph_relationship(rel: GraphRelationship):
|
|
"""Create a relationship between two nodes"""
|
|
if not neo4j_available or not neo4j_driver:
|
|
raise HTTPException(status_code=503, detail="Neo4j not available")
|
|
|
|
try:
|
|
async with neo4j_driver.session() as session:
|
|
props_clause = ""
|
|
if rel.properties:
|
|
props_clause = " SET r += $properties"
|
|
|
|
cypher = f"""
|
|
MATCH (a {{node_id: $from_id}})
|
|
MATCH (b {{node_id: $to_id}})
|
|
MERGE (a)-[r:{rel.relationship_type}]->(b)
|
|
{props_clause}
|
|
SET r.created_at = datetime()
|
|
RETURN a.node_id as from_id, b.node_id as to_id, type(r) as rel_type
|
|
"""
|
|
|
|
result = await session.run(
|
|
cypher,
|
|
from_id=rel.from_node_id,
|
|
to_id=rel.to_node_id,
|
|
properties=rel.properties or {}
|
|
)
|
|
record = await result.single()
|
|
|
|
if record:
|
|
logger.info(f"🔗 Created relationship: {rel.from_node_id} -[{rel.relationship_type}]-> {rel.to_node_id}")
|
|
return {
|
|
"status": "ok",
|
|
"from_id": record["from_id"],
|
|
"to_id": record["to_id"],
|
|
"relationship_type": record["rel_type"]
|
|
}
|
|
|
|
raise HTTPException(status_code=404, detail="One or both nodes not found")
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Neo4j error creating relationship: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/v1/graph/query")
|
|
async def query_graph(query: GraphQuery):
|
|
"""Query the knowledge graph"""
|
|
if not neo4j_available or not neo4j_driver:
|
|
raise HTTPException(status_code=503, detail="Neo4j not available")
|
|
|
|
try:
|
|
async with neo4j_driver.session() as session:
|
|
# If direct Cypher provided, use it (with safety check)
|
|
if query.cypher:
|
|
# Basic safety: only allow read queries
|
|
if any(kw in query.cypher.upper() for kw in ["DELETE", "REMOVE", "DROP", "CREATE", "MERGE", "SET"]):
|
|
raise HTTPException(status_code=400, detail="Only read queries allowed via cypher parameter")
|
|
|
|
result = await session.run(query.cypher)
|
|
records = await result.data()
|
|
return {"status": "ok", "results": records, "count": len(records)}
|
|
|
|
# Build structured query
|
|
if query.node_id:
|
|
# Get specific node with relationships
|
|
cypher = f"""
|
|
MATCH (n {{node_id: $node_id}})
|
|
OPTIONAL MATCH (n)-[r]-(related)
|
|
RETURN n, collect({{rel: type(r), node: related}}) as connections
|
|
LIMIT 1
|
|
"""
|
|
result = await session.run(cypher, node_id=query.node_id)
|
|
|
|
elif query.node_label:
|
|
# Get nodes by label
|
|
cypher = f"""
|
|
MATCH (n:{query.node_label})
|
|
RETURN n
|
|
ORDER BY n.updated_at DESC
|
|
LIMIT $limit
|
|
"""
|
|
result = await session.run(cypher, limit=query.limit)
|
|
|
|
else:
|
|
# Get recent nodes
|
|
cypher = """
|
|
MATCH (n)
|
|
RETURN n, labels(n) as labels
|
|
ORDER BY n.updated_at DESC
|
|
LIMIT $limit
|
|
"""
|
|
result = await session.run(cypher, limit=query.limit)
|
|
|
|
records = await result.data()
|
|
return {"status": "ok", "results": records, "count": len(records)}
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Neo4j query error: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.get("/v1/graph/search")
|
|
async def search_graph(q: str, types: Optional[str] = None, limit: int = 20):
|
|
"""Search nodes by text in properties"""
|
|
if not neo4j_available or not neo4j_driver:
|
|
raise HTTPException(status_code=503, detail="Neo4j not available")
|
|
|
|
try:
|
|
async with neo4j_driver.session() as session:
|
|
# Build label filter
|
|
label_filter = ""
|
|
if types:
|
|
labels = [t.strip() for t in types.split(",")]
|
|
label_filter = " AND (" + " OR ".join([f"n:{l}" for l in labels]) + ")"
|
|
|
|
# Search in common text properties
|
|
cypher = f"""
|
|
MATCH (n)
|
|
WHERE (
|
|
n.name CONTAINS $query OR
|
|
n.title CONTAINS $query OR
|
|
n.text CONTAINS $query OR
|
|
n.description CONTAINS $query OR
|
|
n.content CONTAINS $query
|
|
){label_filter}
|
|
RETURN n, labels(n) as labels
|
|
ORDER BY n.updated_at DESC
|
|
LIMIT $limit
|
|
"""
|
|
|
|
result = await session.run(cypher, query=q, limit=limit)
|
|
records = await result.data()
|
|
|
|
return {"status": "ok", "query": q, "results": records, "count": len(records)}
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Neo4j search error: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.get("/v1/graph/stats")
|
|
async def get_graph_stats():
|
|
"""Get knowledge graph statistics"""
|
|
if not neo4j_available or not neo4j_driver:
|
|
raise HTTPException(status_code=503, detail="Neo4j not available")
|
|
|
|
try:
|
|
async with neo4j_driver.session() as session:
|
|
# Get node counts by label
|
|
labels_result = await session.run("""
|
|
CALL db.labels() YIELD label
|
|
CALL apoc.cypher.run('MATCH (n:`' + label + '`) RETURN count(n) as count', {}) YIELD value
|
|
RETURN label, value.count as count
|
|
""")
|
|
|
|
# If APOC not available, use simpler query
|
|
try:
|
|
labels_data = await labels_result.data()
|
|
except:
|
|
labels_result = await session.run("""
|
|
MATCH (n)
|
|
RETURN labels(n)[0] as label, count(*) as count
|
|
ORDER BY count DESC
|
|
""")
|
|
labels_data = await labels_result.data()
|
|
|
|
# Get relationship counts
|
|
rels_result = await session.run("""
|
|
MATCH ()-[r]->()
|
|
RETURN type(r) as type, count(*) as count
|
|
ORDER BY count DESC
|
|
""")
|
|
rels_data = await rels_result.data()
|
|
|
|
# Get total counts
|
|
total_result = await session.run("""
|
|
MATCH (n) RETURN count(n) as nodes
|
|
""")
|
|
total_nodes = (await total_result.single())["nodes"]
|
|
|
|
total_rels_result = await session.run("""
|
|
MATCH ()-[r]->() RETURN count(r) as relationships
|
|
""")
|
|
total_rels = (await total_rels_result.single())["relationships"]
|
|
|
|
return {
|
|
"status": "ok",
|
|
"total_nodes": total_nodes,
|
|
"total_relationships": total_rels,
|
|
"nodes_by_label": labels_data,
|
|
"relationships_by_type": rels_data
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Neo4j stats error: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.on_event("shutdown")
|
|
async def shutdown_event():
|
|
"""Cleanup connections on shutdown"""
|
|
global neo4j_driver, http_client, nc
|
|
|
|
if neo4j_driver:
|
|
await neo4j_driver.close()
|
|
logger.info("🔌 Neo4j connection closed")
|
|
|
|
if http_client:
|
|
await http_client.aclose()
|
|
logger.info("🔌 HTTP client closed")
|
|
|
|
if nc:
|
|
await nc.close()
|
|
logger.info("🔌 NATS connection closed")
|
|
|