459 lines
15 KiB
Python
459 lines
15 KiB
Python
import os
|
|
import logging
|
|
import json
|
|
import time
|
|
from aiohttp import web
|
|
from aiohttp_cors import setup as cors_setup, ResourceOptions
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables from .env file
|
|
load_dotenv()
|
|
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Import the agent logic
|
|
from src.agents import (
|
|
get_simple_litellm_agent,
|
|
get_pydantic_agent,
|
|
PYDANTIC_AI_AVAILABLE
|
|
)
|
|
from src.agents.multi_stage_agent import create_multi_stage_agent
|
|
from src.agents.stream_handler import get_stream_handler
|
|
from src.tools import get_all_tools
|
|
from src.utils import extract_user_id_from_request
|
|
|
|
def format_status_message(message: str, phase: str) -> str:
|
|
"""
|
|
Format status messages with butler-appropriate box-drawing characters.
|
|
|
|
Uses Unicode box-drawing characters for visual structure:
|
|
- ┌─ for starting messages
|
|
- ├─ for continuing messages
|
|
- └─ for completing messages
|
|
"""
|
|
if phase == "analysis":
|
|
# Starting steward consultation
|
|
return f"┌─ {message}"
|
|
elif phase == "analysis_complete":
|
|
# Steward consultation complete
|
|
return f"└─ {message}"
|
|
elif phase == "tool_execution":
|
|
# Tool being executed
|
|
return f"├─ {message}"
|
|
elif phase == "fallback":
|
|
# Fallback/warning
|
|
return f"└─ {message}"
|
|
else:
|
|
# Default
|
|
return f"├─ {message}"
|
|
|
|
async def chat_completions(request):
|
|
"""
|
|
Handles OpenAI-compatible chat completion requests using PydanticAI.
|
|
Default endpoint - uses PydanticAI with tools enabled.
|
|
"""
|
|
if not PYDANTIC_AI_AVAILABLE:
|
|
return web.json_response({
|
|
"error": {"message": "PydanticAI not available. Install with: pip install pydantic-ai"}
|
|
}, status=503)
|
|
|
|
from src.metrics import get_metrics_collector
|
|
metrics = get_metrics_collector()
|
|
metrics.increment_concurrent_requests()
|
|
start_time = time.time()
|
|
agent_type = "pydantic"
|
|
success = False
|
|
error_msg = None
|
|
|
|
try:
|
|
data = await request.json()
|
|
logger.info(f"[DEFAULT/PYDANTIC_AI] Received chat request")
|
|
|
|
# Extract relevant fields from the request
|
|
messages = data.get("messages")
|
|
model = data.get("model", "Tatlock")
|
|
stream = data.get("stream", False)
|
|
conversation_id = data.get("conversation_id")
|
|
enable_tools = data.get("enable_tools", True)
|
|
multi_stage_analysis = data.get("multi_stage_analysis", True) # Enable by default
|
|
|
|
# Extract user ID from request
|
|
user_id = extract_user_id_from_request(data)
|
|
|
|
if not messages:
|
|
raise web.HTTPBadRequest(reason="'messages' field is required")
|
|
|
|
# Get the agent instance
|
|
base_agent = get_pydantic_agent(discover_tools=enable_tools, user_id=user_id)
|
|
|
|
# Wrap with multi-stage orchestration if enabled
|
|
agent = create_multi_stage_agent(base_agent, enable_multi_stage=multi_stage_analysis)
|
|
|
|
# For non-streaming requests, collect the full response
|
|
if not stream:
|
|
# Collect all content chunks from two-stage agent
|
|
full_content = []
|
|
async for chunk in agent.chat_with_analysis(
|
|
messages=messages,
|
|
conversation_id=conversation_id,
|
|
stream=False
|
|
):
|
|
if chunk.get("type") == "content":
|
|
full_content.append(chunk.get("content", ""))
|
|
|
|
response_content = "".join(full_content)
|
|
success = True
|
|
return web.json_response({
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": response_content
|
|
},
|
|
"finish_reason": "stop"
|
|
}],
|
|
"model": model,
|
|
"usage": {
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"total_tokens": 0
|
|
},
|
|
"tools_enabled": enable_tools,
|
|
"multi_stage_analysis": multi_stage_analysis
|
|
})
|
|
else:
|
|
# Handle streaming response with StreamHandler
|
|
response = web.StreamResponse(
|
|
status=200,
|
|
reason='OK',
|
|
headers={
|
|
'Content-Type': 'text/event-stream',
|
|
'Cache-Control': 'no-cache',
|
|
'Connection': 'keep-alive',
|
|
}
|
|
)
|
|
await response.prepare(request)
|
|
|
|
try:
|
|
# Create stream handler instance
|
|
handler = get_stream_handler(model_name=model)
|
|
|
|
# Process stream through handler
|
|
async for chunk in handler.process_stream(
|
|
messages=messages,
|
|
multi_stage_enabled=multi_stage_analysis
|
|
):
|
|
# Stream handler returns pre-formatted messages
|
|
await response.write(f"data: {json.dumps(chunk)}\n\n".encode('utf-8'))
|
|
|
|
await response.write(b"data: [DONE]\n\n")
|
|
success = True
|
|
finally:
|
|
await response.write_eof()
|
|
|
|
return response
|
|
|
|
except web.HTTPBadRequest:
|
|
raise
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
logger.exception(f"Error in chat_completions: {e}")
|
|
return web.json_response({
|
|
"error": {"message": f"Internal server error: {str(e)}"}
|
|
}, status=500)
|
|
finally:
|
|
duration_ms = (time.time() - start_time) * 1000
|
|
metrics.decrement_concurrent_requests()
|
|
metrics.record_request(
|
|
agent_type=agent_type,
|
|
duration_ms=duration_ms,
|
|
success=success,
|
|
streaming=data.get("stream", False) if 'data' in locals() else False,
|
|
user_id=user_id if 'user_id' in locals() else None,
|
|
error=error_msg
|
|
)
|
|
|
|
async def chat_simple(request):
|
|
"""
|
|
Handles chat requests using SimpleLiteLLMAgent (fallback, no tools).
|
|
Endpoint: /v1/chat/simple
|
|
"""
|
|
from src.metrics import get_metrics_collector
|
|
metrics = get_metrics_collector()
|
|
metrics.increment_concurrent_requests()
|
|
start_time = time.time()
|
|
agent_type = "simple"
|
|
success = False
|
|
error_msg = None
|
|
|
|
try:
|
|
data = await request.json()
|
|
logger.info(f"[SIMPLE/LITELLM] Received chat request")
|
|
|
|
# Extract relevant fields from the request
|
|
messages = data.get("messages")
|
|
model = data.get("model", "simple")
|
|
stream = data.get("stream", False)
|
|
conversation_id = data.get("conversation_id")
|
|
user_id = extract_user_id_from_request(data)
|
|
|
|
if not messages:
|
|
raise web.HTTPBadRequest(reason="'messages' field is required")
|
|
|
|
# Get simple agent instance (no tools)
|
|
agent = get_simple_litellm_agent()
|
|
|
|
# For non-streaming requests
|
|
if not stream:
|
|
response_content = await agent.chat_completion(
|
|
messages=messages,
|
|
conversation_id=conversation_id
|
|
)
|
|
return web.json_response({
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": response_content
|
|
},
|
|
"finish_reason": "stop"
|
|
}],
|
|
"model": model,
|
|
"usage": {
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"total_tokens": 0
|
|
}
|
|
})
|
|
else:
|
|
# Streaming response
|
|
response = web.StreamResponse(
|
|
status=200,
|
|
reason='OK',
|
|
headers={
|
|
'Content-Type': 'text/event-stream',
|
|
'Cache-Control': 'no-cache',
|
|
'Connection': 'keep-alive',
|
|
}
|
|
)
|
|
await response.prepare(request)
|
|
|
|
try:
|
|
async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
|
|
if chunk["type"] == "content":
|
|
chunk_data = {
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"content": chunk["content"]},
|
|
"finish_reason": chunk.get("finish_reason")
|
|
}],
|
|
"model": model
|
|
}
|
|
await response.write(f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8'))
|
|
|
|
await response.write(b"data: [DONE]\n\n")
|
|
success = True
|
|
finally:
|
|
await response.write_eof()
|
|
|
|
return response
|
|
|
|
except web.HTTPBadRequest:
|
|
raise
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
logger.exception(f"Error in chat_simple: {e}")
|
|
return web.json_response({
|
|
"error": {"message": f"Internal server error: {str(e)}"}
|
|
}, status=500)
|
|
finally:
|
|
duration_ms = (time.time() - start_time) * 1000
|
|
metrics.decrement_concurrent_requests()
|
|
metrics.record_request(
|
|
agent_type=agent_type,
|
|
duration_ms=duration_ms,
|
|
success=success,
|
|
streaming=data.get("stream", False) if 'data' in locals() else False,
|
|
user_id=user_id if 'user_id' in locals() else None,
|
|
error=error_msg
|
|
)
|
|
|
|
async def list_models(request):
|
|
"""List available models"""
|
|
return web.json_response({
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"id": "Tatlock",
|
|
"object": "model",
|
|
"created": int(time.time()),
|
|
"owned_by": "core-ai",
|
|
"description": "PydanticAI agent with full tool support - your British butler assistant"
|
|
}
|
|
]
|
|
})
|
|
|
|
async def list_tools(request):
|
|
"""List all available tools"""
|
|
try:
|
|
tools = get_all_tools()
|
|
tool_list = []
|
|
|
|
for name, func in tools.items():
|
|
import inspect
|
|
doc = inspect.getdoc(func) or "No description"
|
|
tool_list.append({
|
|
"name": name,
|
|
"description": doc.split('\n')[0],
|
|
"type": "local" if not name.startswith("core-api__") else "openapi"
|
|
})
|
|
|
|
return web.json_response({
|
|
"tools": tool_list,
|
|
"tools_count": len(tool_list)
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Error listing tools: {e}")
|
|
return web.json_response({
|
|
"error": {"message": str(e)}
|
|
}, status=500)
|
|
|
|
async def health_check(request):
|
|
"""Health check endpoint"""
|
|
return web.json_response({
|
|
"status": "ok",
|
|
"service": "core-ai",
|
|
"agents": {
|
|
"Tatlock": PYDANTIC_AI_AVAILABLE
|
|
},
|
|
"default_agent": "Tatlock" if PYDANTIC_AI_AVAILABLE else None,
|
|
"tools_count": len(get_all_tools())
|
|
})
|
|
|
|
async def get_metrics(request):
|
|
"""
|
|
Get comprehensive performance metrics.
|
|
|
|
Returns detailed statistics on agent performance, tool execution,
|
|
memory system, and request patterns.
|
|
"""
|
|
try:
|
|
from src.metrics import get_metrics_collector
|
|
metrics_collector = get_metrics_collector()
|
|
metrics = metrics_collector.get_metrics()
|
|
return web.json_response(metrics)
|
|
except Exception as e:
|
|
logger.exception(f"Error getting metrics: {e}")
|
|
return web.json_response({
|
|
"error": {"message": str(e)}
|
|
}, status=500)
|
|
|
|
async def get_recent_errors(request):
|
|
"""Get recent request errors."""
|
|
try:
|
|
from src.metrics import get_metrics_collector
|
|
metrics_collector = get_metrics_collector()
|
|
|
|
limit = int(request.query.get('limit', 20))
|
|
errors = metrics_collector.get_recent_errors(limit=limit)
|
|
|
|
return web.json_response({
|
|
"errors": errors,
|
|
"total": len(errors)
|
|
})
|
|
except Exception as e:
|
|
logger.exception(f"Error getting errors: {e}")
|
|
return web.json_response({
|
|
"error": {"message": str(e)}
|
|
}, status=500)
|
|
|
|
async def get_tool_failures(request):
|
|
"""Get recent tool execution failures."""
|
|
try:
|
|
from src.metrics import get_metrics_collector
|
|
metrics_collector = get_metrics_collector()
|
|
|
|
limit = int(request.query.get('limit', 20))
|
|
failures = metrics_collector.get_recent_tool_failures(limit=limit)
|
|
|
|
return web.json_response({
|
|
"failures": failures,
|
|
"total": len(failures)
|
|
})
|
|
except Exception as e:
|
|
logger.exception(f"Error getting tool failures: {e}")
|
|
return web.json_response({
|
|
"error": {"message": str(e)}
|
|
}, status=500)
|
|
|
|
async def reset_metrics(request):
|
|
"""Reset all metrics (admin endpoint)."""
|
|
try:
|
|
from src.metrics import get_metrics_collector
|
|
metrics_collector = get_metrics_collector()
|
|
metrics_collector.reset()
|
|
|
|
return web.json_response({
|
|
"message": "Metrics reset successfully"
|
|
})
|
|
except Exception as e:
|
|
logger.exception(f"Error resetting metrics: {e}")
|
|
return web.json_response({
|
|
"error": {"message": str(e)}
|
|
}, status=500)
|
|
|
|
async def setup_routes(app):
|
|
# Chat endpoints
|
|
app.router.add_post("/chat/completions", chat_completions) # Alias without /v1
|
|
app.router.add_post("/v1/chat/completions", chat_completions) # Default: PydanticAI
|
|
app.router.add_post("/v1/chat/simple", chat_simple) # Fallback: Simple agent
|
|
|
|
# Models endpoint
|
|
app.router.add_get("/models", list_models)
|
|
app.router.add_get("/v1/models", list_models)
|
|
|
|
# Tools endpoint
|
|
app.router.add_get("/tools", list_tools)
|
|
app.router.add_get("/v1/tools", list_tools)
|
|
|
|
# Health check
|
|
app.router.add_get("/health", health_check)
|
|
|
|
# Metrics endpoints
|
|
app.router.add_get("/metrics", get_metrics)
|
|
app.router.add_get("/metrics/errors", get_recent_errors)
|
|
app.router.add_get("/metrics/tool-failures", get_tool_failures)
|
|
app.router.add_post("/metrics/reset", reset_metrics)
|
|
|
|
# Setup CORS
|
|
cors = cors_setup(app, defaults={
|
|
"*": ResourceOptions(
|
|
allow_credentials=True,
|
|
expose_headers="*",
|
|
allow_headers="*",
|
|
allow_methods="*"
|
|
)
|
|
})
|
|
|
|
# Apply CORS to all routes
|
|
for route in list(app.router.routes()):
|
|
cors.add(route)
|
|
|
|
def main():
|
|
app = web.Application()
|
|
|
|
# Set up routes
|
|
import asyncio
|
|
loop = asyncio.get_event_loop()
|
|
loop.run_until_complete(setup_routes(app))
|
|
|
|
# Run the application
|
|
logger.info("Starting core-ai service on http://0.0.0.0:8086")
|
|
web.run_app(app, host='0.0.0.0', port=8086)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|