jpmschweitzer
8487b2a366
Introduces a comprehensive, multi-tiered memory system to provide conversation history and context for the AI agent. This lays the foundation for more stateful and intelligent interactions. Key components of this implementation: - **Multi-Tiered Memory Architecture:** - **Tier 1 (Working Memory):** A fast, in-memory buffer (`ConversationBufferMemory`) that holds the most recent turns of a conversation for immediate access. - **Tier 3 (Long-Term Memory):** A persistent, semantic search-based memory store using Qdrant (`QdrantConversationMemory`). It stores all conversation turns as vector embeddings, enabling long-term recall and similarity search. - **Qdrant Integration:** - The `qdrant-client` is added to manage collections and perform vector search operations. - Each user is assigned a dedicated Qdrant collection for multi-tenancy. - **Ollama Embedding Client:** - A new `OllamaEmbeddingClient` generates text embeddings via the Ollama API, replacing the need for local sentence-transformer models. This significantly reduces the service's dependency footprint. - **Configuration and Stack Updates:** - The `config.py` and `core-ai.yml` stack file are updated with new settings for enabling memory, configuring Qdrant, and specifying the embedding model. - **Utility and Schema Additions:** - New Pydantic schemas (`memory/schemas.py`) define the data structures for conversation turns and memory management. - Utility functions (`utils.py`) are added for user ID sanitization and collection naming. This feature enhances the agent's capabilities by allowing it to maintain context across multiple turns and sessions, leading to more coherent and relevant responses.
417 lines
15 KiB
Python
417 lines
15 KiB
Python
import os
|
|
import logging
|
|
import json
|
|
import time # Import time module
|
|
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 agent.
|
|
Default endpoint - uses PydanticAI Agent 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)
|
|
|
|
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) # Tools enabled by default
|
|
|
|
# Extract user ID from request (supports user_id, user_email, or falls back to default)
|
|
user_id = extract_user_id_from_request(data)
|
|
|
|
if not messages:
|
|
raise web.HTTPBadRequest(reason="'messages' field is required")
|
|
|
|
# Get the agent instance (default: PydanticAI agent with tools)
|
|
# Note: Agent is cached per user_id, so each user gets their own agent instance with their memory
|
|
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
|
|
)
|
|
return web.json_response({
|
|
"id": f"chatcmpl-{os.urandom(12).hex()}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": "pydantic",
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": response_content},
|
|
"finish_reason": "stop"
|
|
}],
|
|
"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,
|
|
headers={'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive'}
|
|
)
|
|
await response.prepare(request)
|
|
|
|
async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
|
|
chunk_type = chunk.get("type", "content")
|
|
|
|
if chunk_type == "content":
|
|
json_chunk = {
|
|
"id": f"chatcmpl-{os.urandom(12).hex()}",
|
|
"object": "chat.completion.chunk",
|
|
"created": int(time.time()),
|
|
"model": "pydantic",
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"content": chunk.get("content", "")},
|
|
"finish_reason": chunk.get("finish_reason")
|
|
}]
|
|
}
|
|
await response.write(f"data: {json.dumps(json_chunk)}\n\n".encode())
|
|
|
|
if chunk.get("finish_reason") == "stop":
|
|
break
|
|
elif chunk_type == "error":
|
|
error_chunk = {
|
|
"error": {"message": chunk.get("content", "Unknown error")}
|
|
}
|
|
await response.write(f"data: {json.dumps(error_chunk)}\n\n".encode())
|
|
break
|
|
|
|
await response.write(b"data: [DONE]\n\n")
|
|
await response.write_eof()
|
|
return response
|
|
|
|
except web.HTTPBadRequest as e:
|
|
logger.warning(f"Bad request: {e.reason}")
|
|
return web.json_response({"error": {"message": e.reason}}, status=400)
|
|
except Exception as e:
|
|
logger.exception("[DEFAULT/PYDANTIC_AI] Error during chat completion:")
|
|
return web.json_response({"error": {"message": str(e)}}, status=500)
|
|
|
|
async def chat_simple(request):
|
|
"""
|
|
Handles chat requests using SimpleLiteLLMAgent (no tools).
|
|
Endpoint: POST /v1/chat/simple
|
|
"""
|
|
try:
|
|
data = await request.json()
|
|
logger.info(f"[SIMPLE] Received chat request")
|
|
|
|
messages = data.get("messages")
|
|
model = data.get("model", "simple")
|
|
stream = data.get("stream", False)
|
|
conversation_id = data.get("conversation_id")
|
|
|
|
if not messages:
|
|
raise web.HTTPBadRequest(reason="'messages' field is required")
|
|
|
|
# Get SimpleLiteLLM agent
|
|
agent = get_simple_litellm_agent()
|
|
|
|
# Non-streaming response
|
|
if not stream:
|
|
response_content = await agent.chat_completion(
|
|
messages=messages,
|
|
conversation_id=conversation_id
|
|
)
|
|
return web.json_response({
|
|
"id": f"chatcmpl-{os.urandom(12).hex()}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": "simple",
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": response_content},
|
|
"finish_reason": "stop"
|
|
}],
|
|
"usage": {
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"total_tokens": 0
|
|
}
|
|
})
|
|
else:
|
|
# Streaming response
|
|
response = web.StreamResponse(
|
|
status=200,
|
|
headers={'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive'}
|
|
)
|
|
await response.prepare(request)
|
|
|
|
async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
|
|
json_chunk = {
|
|
"id": f"chatcmpl-{os.urandom(12).hex()}",
|
|
"object": "chat.completion.chunk",
|
|
"created": int(time.time()),
|
|
"model": "simple",
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"content": chunk.get("content", "")},
|
|
"finish_reason": chunk.get("finish_reason")
|
|
}]
|
|
}
|
|
await response.write(f"data: {json.dumps(json_chunk)}\n\n".encode())
|
|
if chunk.get("finish_reason") == "stop":
|
|
break
|
|
|
|
await response.write(b"data: [DONE]\n\n")
|
|
await response.write_eof()
|
|
return response
|
|
|
|
except web.HTTPBadRequest as e:
|
|
logger.warning(f"Bad request: {e.reason}")
|
|
return web.json_response({"error": {"message": e.reason}}, status=400)
|
|
except Exception as e:
|
|
logger.exception("[SIMPLE] Error during chat completion:")
|
|
return web.json_response({"error": {"message": str(e)}}, status=500)
|
|
|
|
|
|
async def chat_pydantic(request):
|
|
"""
|
|
Handles chat requests using PydanticAI Agent with tools.
|
|
Endpoint: POST /v1/chat/pydantic
|
|
"""
|
|
if not PYDANTIC_AI_AVAILABLE:
|
|
return web.json_response({
|
|
"error": {"message": "PydanticAI not available. Install with: pip install pydantic-ai"}
|
|
}, status=503)
|
|
|
|
try:
|
|
data = await request.json()
|
|
logger.info(f"[PYDANTIC_AI] Received chat 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 PydanticAI agent with or without tools
|
|
agent = get_pydantic_agent(discover_tools=enable_tools, user_id=user_id)
|
|
|
|
# Non-streaming response
|
|
if not stream:
|
|
response_content = await agent.chat_completion(
|
|
messages=messages,
|
|
conversation_id=conversation_id
|
|
)
|
|
return web.json_response({
|
|
"id": f"chatcmpl-{os.urandom(12).hex()}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": "pydantic",
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": response_content},
|
|
"finish_reason": "stop"
|
|
}],
|
|
"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:
|
|
# Streaming response
|
|
response = web.StreamResponse(
|
|
status=200,
|
|
headers={'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive'}
|
|
)
|
|
await response.prepare(request)
|
|
|
|
async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
|
|
chunk_type = chunk.get("type", "content")
|
|
|
|
if chunk_type == "content":
|
|
json_chunk = {
|
|
"id": f"chatcmpl-{os.urandom(12).hex()}",
|
|
"object": "chat.completion.chunk",
|
|
"created": int(time.time()),
|
|
"model": "pydantic",
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"content": chunk.get("content", "")},
|
|
"finish_reason": chunk.get("finish_reason")
|
|
}]
|
|
}
|
|
await response.write(f"data: {json.dumps(json_chunk)}\n\n".encode())
|
|
|
|
if chunk.get("finish_reason") == "stop":
|
|
break
|
|
elif chunk_type == "error":
|
|
error_chunk = {
|
|
"error": {"message": chunk.get("content", "Unknown error")}
|
|
}
|
|
await response.write(f"data: {json.dumps(error_chunk)}\n\n".encode())
|
|
break
|
|
|
|
await response.write(b"data: [DONE]\n\n")
|
|
await response.write_eof()
|
|
return response
|
|
|
|
except web.HTTPBadRequest as e:
|
|
logger.warning(f"Bad request: {e.reason}")
|
|
return web.json_response({"error": {"message": e.reason}}, status=400)
|
|
except Exception as e:
|
|
logger.exception("[PYDANTIC_AI] Error during chat completion:")
|
|
return web.json_response({"error": {"message": str(e)}}, status=500)
|
|
|
|
|
|
async def list_models(request):
|
|
"""
|
|
Lists available models (OpenAI-compatible endpoint).
|
|
Endpoint: GET /v1/models
|
|
"""
|
|
models = [
|
|
{
|
|
"id": "Tatlock",
|
|
"object": "model",
|
|
"created": int(time.time()),
|
|
"owned_by": "core-ai",
|
|
"permission": [],
|
|
"root": "tatlock",
|
|
"parent": None,
|
|
},
|
|
{
|
|
"id": "simple",
|
|
"object": "model",
|
|
"created": int(time.time()),
|
|
"owned_by": "core-ai",
|
|
"permission": [],
|
|
"root": "simple",
|
|
"parent": None,
|
|
}
|
|
]
|
|
|
|
return web.json_response({
|
|
"object": "list",
|
|
"data": models
|
|
})
|
|
|
|
|
|
async def list_tools(request):
|
|
"""
|
|
Lists all available tools.
|
|
Endpoint: GET /v1/tools
|
|
"""
|
|
try:
|
|
tools = get_all_tools()
|
|
|
|
tools_info = []
|
|
for name, func in tools.items():
|
|
tools_info.append({
|
|
"name": name,
|
|
"description": func.__doc__.strip() if func.__doc__ else "No description available",
|
|
"type": "local"
|
|
})
|
|
|
|
return web.json_response({
|
|
"tools": tools_info,
|
|
"count": len(tools_info),
|
|
"pydantic_ai_available": PYDANTIC_AI_AVAILABLE
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.exception("Error listing tools:")
|
|
return web.json_response({"error": {"message": str(e)}}, status=500)
|
|
|
|
|
|
async def health_check(request):
|
|
"""Simple 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 setup_routes(app):
|
|
# Chat endpoints
|
|
app.router.add_post("/chat/completions", chat_completions) # Alias without /v1 for compatibility
|
|
app.router.add_post("/v1/chat/completions", chat_completions) # Default (PydanticAI)
|
|
app.router.add_post("/v1/chat/simple", chat_simple) # Simple agent (no tools)
|
|
app.router.add_post("/v1/chat/pydantic", chat_pydantic) # Alias for default
|
|
|
|
# OpenAI-compatible endpoints
|
|
app.router.add_get("/v1/models", list_models) # List available models
|
|
app.router.add_get("/models", list_models) # Alias without /v1 prefix
|
|
|
|
# Tool management
|
|
app.router.add_get("/v1/tools", list_tools) # List available tools
|
|
|
|
# Health check
|
|
app.router.add_get("/health", health_check)
|
|
|
|
# Setup CORS
|
|
cors = cors_setup(app, defaults={
|
|
"*": ResourceOptions(
|
|
allow_credentials=True,
|
|
expose_headers="*",
|
|
allow_headers="*",
|
|
allow_methods="*"
|
|
)
|
|
})
|
|
|
|
# Configure CORS on all routes
|
|
for route in list(app.router.routes()):
|
|
cors.add(route)
|
|
|
|
def main():
|
|
app = web.Application()
|
|
app.on_startup.append(setup_routes) # Register routes on startup
|
|
|
|
# Configuration
|
|
host = os.getenv("HOST", "0.0.0.0")
|
|
port = int(os.getenv("PORT", 8086)) # Use 8086 to avoid conflict with core-ai
|
|
|
|
logger.info(f"Starting core-ai service on http://{host}:{port}")
|
|
web.run_app(app, host=host, port=port)
|
|
|
|
if __name__ == "__main__":
|
|
main() |