jpmschweitzerandClaude
632b20febe
Adds comprehensive in-memory metrics collection for monitoring AI agent performance, tool execution, and system behavior. New Components: - src/metrics/collector.py: Thread-safe MetricsCollector class - Tracks agent requests (response times, errors, concurrency) - Tracks tool execution (calls, success/failure, durations) - Tracks memory system (tier1/tier2 hits, consolidations) - Calculates percentiles (p50, p95, p99) for performance analysis - Sliding window retention (1h detailed, 24h aggregated) - src/metrics/decorators.py: Automatic instrumentation decorators - @track_tool_execution: Auto-tracks tool calls with metrics - @track_duration: Generic duration tracking decorator - src/metrics/__init__.py: Module exports API Endpoints: - GET /metrics: Comprehensive performance metrics snapshot - GET /metrics/errors: Recent request errors with timestamps - GET /metrics/tool-failures: Recent tool execution failures - POST /metrics/reset: Clear all metrics (admin endpoint) Instrumentation: - Enhanced main.py chat handlers with metrics tracking - Modified tools/registry.py log_tool_call to track execution metrics - All metrics recorded with proper error handling and context Features: - Thread-safe with threading.Lock for concurrent requests - No database dependencies (in-memory only) - Automatic cleanup of old data (sliding windows) - Detailed statistics: avg, p50, p95, p99 response times - Per-user tracking and request attribution - Tool success rates and performance analysis Tested and validated: - All endpoints responding correctly - Request metrics collected successfully - Response time percentiles calculated correctly - User tracking functional 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
428 lines
14 KiB
Python
428 lines
14 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.tools import get_all_tools
|
|
from src.utils import extract_user_id_from_request
|
|
|
|
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", "pydantic")
|
|
stream = data.get("stream", False)
|
|
conversation_id = data.get("conversation_id")
|
|
enable_tools = data.get("enable_tools", True)
|
|
|
|
# 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
|
|
agent = get_pydantic_agent(discover_tools=enable_tools, user_id=user_id)
|
|
|
|
# For non-streaming requests, collect the full response
|
|
if not stream:
|
|
response_content = await agent.chat_completion(
|
|
messages=messages,
|
|
conversation_id=conversation_id
|
|
)
|
|
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,
|
|
"tools_count": len(agent.tools) if enable_tools else 0
|
|
})
|
|
else:
|
|
# Handle 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_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": "pydantic",
|
|
"object": "model",
|
|
"created": int(time.time()),
|
|
"owned_by": "core-ai"
|
|
},
|
|
{
|
|
"id": "simple",
|
|
"object": "model",
|
|
"created": int(time.time()),
|
|
"owned_by": "core-ai"
|
|
}
|
|
]
|
|
})
|
|
|
|
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": {
|
|
"simple": True,
|
|
"pydantic": PYDANTIC_AI_AVAILABLE
|
|
},
|
|
"default_agent": "pydantic" if PYDANTIC_AI_AVAILABLE else "simple",
|
|
"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()
|