ok... ok... I'll add it to git...
This commit is contained in:
@@ -0,0 +1,6 @@
|
||||
"""
|
||||
Core Code API - OpenAPI-compatible functions for Open WebUI
|
||||
"""
|
||||
|
||||
__version__ = "1.0.0"
|
||||
__author__ = "Core Code Team"
|
||||
@@ -0,0 +1,291 @@
|
||||
"""
|
||||
OpenAI-compatible /v1/chat/completions endpoint
|
||||
|
||||
Phase 2: Integrated with memory system for conversation persistence.
|
||||
"""
|
||||
|
||||
import time
|
||||
import logging
|
||||
import uuid
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from typing import AsyncIterator
|
||||
|
||||
from .schemas import (
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionChoice,
|
||||
ChatMessageResponse,
|
||||
UsageInfo,
|
||||
ChatCompletionStreamResponse,
|
||||
ChatCompletionStreamChoice,
|
||||
DeltaMessage,
|
||||
)
|
||||
from src.models.ollama_client import get_ollama_client
|
||||
from src.memory import get_memory_manager, MessageRole as MemoryMessageRole, TokenUsage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
def build_prompt_from_messages(messages: list) -> str:
|
||||
"""
|
||||
Convert message list to a prompt string.
|
||||
|
||||
In Phase 1, we do simple concatenation.
|
||||
Phase 2 will add proper memory management.
|
||||
"""
|
||||
prompt_parts = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.role.value if hasattr(msg.role, 'value') else msg.role
|
||||
content = msg.content
|
||||
|
||||
if role == "system":
|
||||
prompt_parts.append(f"System: {content}")
|
||||
elif role == "user":
|
||||
prompt_parts.append(f"User: {content}")
|
||||
elif role == "assistant":
|
||||
prompt_parts.append(f"Assistant: {content}")
|
||||
|
||||
prompt_parts.append("Assistant:")
|
||||
return "\n\n".join(prompt_parts)
|
||||
|
||||
|
||||
async def stream_chat_completion(
|
||||
request_id: str,
|
||||
model: str,
|
||||
prompt: str,
|
||||
temperature: float,
|
||||
max_tokens: int | None
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream chat completion in OpenAI SSE format.
|
||||
|
||||
Yields:
|
||||
Server-Sent Events formatted strings
|
||||
"""
|
||||
created = int(time.time())
|
||||
ollama_client = get_ollama_client()
|
||||
|
||||
# First chunk with role
|
||||
first_chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(role="assistant"),
|
||||
finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {first_chunk.model_dump_json()}\n\n"
|
||||
|
||||
# Stream tokens
|
||||
try:
|
||||
async for token in ollama_client.generate_streaming(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens
|
||||
):
|
||||
chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(content=token),
|
||||
finish_reason=None
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Streaming error: {e}")
|
||||
# Send error in OpenAI format
|
||||
import json
|
||||
error_chunk = {
|
||||
"error": {
|
||||
"message": str(e),
|
||||
"type": "server_error"
|
||||
}
|
||||
}
|
||||
yield f"data: {json.dumps(error_chunk)}\n\n"
|
||||
return
|
||||
|
||||
# Final chunk
|
||||
final_chunk = ChatCompletionStreamResponse(
|
||||
id=request_id,
|
||||
created=created,
|
||||
model=model,
|
||||
choices=[
|
||||
ChatCompletionStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(),
|
||||
finish_reason="stop"
|
||||
)
|
||||
]
|
||||
)
|
||||
yield f"data: {final_chunk.model_dump_json()}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
|
||||
async def store_conversation_turn(
|
||||
conversation_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
tokens: dict = None
|
||||
):
|
||||
"""
|
||||
Store a conversation turn in memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
role: Message role (user, assistant, system)
|
||||
content: Message content
|
||||
tokens: Optional token usage dict
|
||||
"""
|
||||
try:
|
||||
memory_manager = get_memory_manager()
|
||||
|
||||
# Convert role string to MemoryMessageRole
|
||||
if role == "user":
|
||||
memory_role = MemoryMessageRole.USER
|
||||
elif role == "assistant":
|
||||
memory_role = MemoryMessageRole.ASSISTANT
|
||||
elif role == "system":
|
||||
memory_role = MemoryMessageRole.SYSTEM
|
||||
else:
|
||||
memory_role = MemoryMessageRole.USER # Default fallback
|
||||
|
||||
# Create TokenUsage if provided
|
||||
token_usage = None
|
||||
if tokens:
|
||||
token_usage = TokenUsage(
|
||||
prompt=tokens.get("prompt", 0),
|
||||
completion=tokens.get("completion", 0),
|
||||
total=tokens.get("total", 0)
|
||||
)
|
||||
|
||||
# Store in memory
|
||||
await memory_manager.add_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=memory_role,
|
||||
content=content,
|
||||
tokens=token_usage
|
||||
)
|
||||
|
||||
logger.debug(f"Stored {role} turn in memory for conversation {conversation_id}")
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the request
|
||||
logger.error(f"Failed to store turn in memory: {e}")
|
||||
|
||||
|
||||
@router.post("/v1/chat/completions")
|
||||
async def chat_completions(request: ChatCompletionRequest):
|
||||
"""
|
||||
OpenAI-compatible chat completions endpoint.
|
||||
Supports both streaming and non-streaming.
|
||||
|
||||
Phase 2: Automatically stores conversations in memory system.
|
||||
"""
|
||||
request_id = f"chatcmpl-{int(time.time() * 1000)}"
|
||||
|
||||
# Generate or use provided conversation_id
|
||||
conversation_id = request.conversation_id or f"conv_{uuid.uuid4().hex[:16]}"
|
||||
|
||||
logger.info(
|
||||
f"Chat request: id={request_id}, model={request.model}, "
|
||||
f"messages={len(request.messages)}, stream={request.stream}, "
|
||||
f"conversation_id={conversation_id}, store_in_memory={request.store_in_memory}"
|
||||
)
|
||||
|
||||
# Store user messages in memory (if enabled)
|
||||
if request.store_in_memory:
|
||||
for msg in request.messages:
|
||||
role = msg.role.value if hasattr(msg.role, 'value') else msg.role
|
||||
if role == "user": # Store latest user message
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role=role,
|
||||
content=msg.content
|
||||
)
|
||||
|
||||
# Build prompt from messages
|
||||
prompt = build_prompt_from_messages(request.messages)
|
||||
|
||||
# Streaming response
|
||||
if request.stream:
|
||||
return StreamingResponse(
|
||||
stream_chat_completion(
|
||||
request_id=request_id,
|
||||
model=request.model,
|
||||
prompt=prompt,
|
||||
temperature=request.temperature,
|
||||
max_tokens=request.max_tokens
|
||||
),
|
||||
media_type="text/event-stream"
|
||||
)
|
||||
|
||||
# Non-streaming response
|
||||
try:
|
||||
ollama_client = get_ollama_client()
|
||||
result = await ollama_client.generate_non_streaming(
|
||||
model=request.model,
|
||||
prompt=prompt,
|
||||
temperature=request.temperature,
|
||||
max_tokens=request.max_tokens
|
||||
)
|
||||
|
||||
assistant_content = result["response"]
|
||||
|
||||
# Store assistant response in memory (if enabled)
|
||||
if request.store_in_memory:
|
||||
await store_conversation_turn(
|
||||
conversation_id=conversation_id,
|
||||
role="assistant",
|
||||
content=assistant_content,
|
||||
tokens=result["tokens"]
|
||||
)
|
||||
|
||||
response = ChatCompletionResponse(
|
||||
id=request_id,
|
||||
created=int(time.time()),
|
||||
model=request.model,
|
||||
choices=[
|
||||
ChatCompletionChoice(
|
||||
index=0,
|
||||
message=ChatMessageResponse(
|
||||
role="assistant",
|
||||
content=assistant_content
|
||||
),
|
||||
finish_reason="stop"
|
||||
)
|
||||
],
|
||||
usage=UsageInfo(
|
||||
prompt_tokens=result["tokens"]["prompt"],
|
||||
completion_tokens=result["tokens"]["completion"],
|
||||
total_tokens=result["tokens"]["total"]
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Chat response: id={request_id}, "
|
||||
f"tokens={result['tokens']['total']}, "
|
||||
f"conversation_id={conversation_id}"
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Chat completion error: {e}")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to generate completion: {str(e)}"
|
||||
)
|
||||
@@ -0,0 +1,338 @@
|
||||
"""
|
||||
Conversation History API Endpoints
|
||||
|
||||
Provides endpoints for managing and querying conversation memory:
|
||||
- List conversations
|
||||
- Get conversation history
|
||||
- Search conversations semantically
|
||||
- Delete conversations
|
||||
"""
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from typing import List, Optional
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from src.memory import get_memory_manager, MessageRole
|
||||
|
||||
router = APIRouter(prefix="/v1/conversations", tags=["conversations"])
|
||||
|
||||
|
||||
# Request/Response Models
|
||||
class SearchRequest(BaseModel):
|
||||
"""Request model for semantic search"""
|
||||
query: str = Field(..., description="Search query")
|
||||
limit: int = Field(5, ge=1, le=50, description="Maximum number of results")
|
||||
|
||||
|
||||
class ConversationTurnResponse(BaseModel):
|
||||
"""Response model for a conversation turn"""
|
||||
turn_number: int
|
||||
role: str
|
||||
content: str
|
||||
timestamp: str
|
||||
tokens_prompt: Optional[int] = None
|
||||
tokens_completion: Optional[int] = None
|
||||
tokens_total: Optional[int] = None
|
||||
metadata: dict = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ConversationHistoryResponse(BaseModel):
|
||||
"""Response model for conversation history"""
|
||||
conversation_id: str
|
||||
turn_count: int
|
||||
total_tokens: int
|
||||
turns: List[ConversationTurnResponse]
|
||||
|
||||
|
||||
class SearchResultResponse(BaseModel):
|
||||
"""Response model for a single search result"""
|
||||
conversation_id: str
|
||||
turn_number: int
|
||||
role: str
|
||||
content: str
|
||||
timestamp: str
|
||||
score: float
|
||||
|
||||
|
||||
class SearchResponse(BaseModel):
|
||||
"""Response model for search results"""
|
||||
query: str
|
||||
results: List[SearchResultResponse]
|
||||
count: int
|
||||
|
||||
|
||||
class ConversationStatsResponse(BaseModel):
|
||||
"""Response model for conversation statistics"""
|
||||
conversation_id: str
|
||||
buffer_turns: int
|
||||
buffer_tokens: int
|
||||
qdrant_turns: int
|
||||
qdrant_tokens: int
|
||||
exists_in_buffer: bool
|
||||
exists_in_qdrant: bool
|
||||
|
||||
|
||||
class DeleteResponse(BaseModel):
|
||||
"""Response model for delete operation"""
|
||||
conversation_id: str
|
||||
deleted: bool
|
||||
message: str
|
||||
|
||||
|
||||
# Endpoints
|
||||
|
||||
@router.get(
|
||||
"/{conversation_id}",
|
||||
response_model=ConversationHistoryResponse,
|
||||
summary="Get conversation history",
|
||||
description="Retrieve complete conversation history including all turns"
|
||||
)
|
||||
async def get_conversation(
|
||||
conversation_id: str,
|
||||
include_buffer: bool = Query(
|
||||
True,
|
||||
description="Include recent turns from buffer that haven't been consolidated yet"
|
||||
)
|
||||
):
|
||||
"""
|
||||
Get complete conversation history
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
include_buffer: Include recent buffer turns not yet consolidated
|
||||
|
||||
Returns:
|
||||
Complete conversation history with all turns
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Get full history
|
||||
turns = await manager.get_full_history(conversation_id, include_buffer=include_buffer)
|
||||
|
||||
if not turns:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Conversation {conversation_id} not found"
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
turn_responses = []
|
||||
total_tokens = 0
|
||||
|
||||
for turn in turns:
|
||||
turn_response = ConversationTurnResponse(
|
||||
turn_number=turn.turn_number,
|
||||
role=turn.role.value if isinstance(turn.role, MessageRole) else turn.role,
|
||||
content=turn.content,
|
||||
timestamp=turn.timestamp.isoformat(),
|
||||
metadata=turn.metadata
|
||||
)
|
||||
|
||||
if turn.tokens:
|
||||
turn_response.tokens_prompt = turn.tokens.prompt
|
||||
turn_response.tokens_completion = turn.tokens.completion
|
||||
turn_response.tokens_total = turn.tokens.total
|
||||
total_tokens += turn.tokens.total
|
||||
|
||||
turn_responses.append(turn_response)
|
||||
|
||||
return ConversationHistoryResponse(
|
||||
conversation_id=conversation_id,
|
||||
turn_count=len(turns),
|
||||
total_tokens=total_tokens,
|
||||
turns=turn_responses
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{conversation_id}/stats",
|
||||
response_model=ConversationStatsResponse,
|
||||
summary="Get conversation statistics",
|
||||
description="Get detailed statistics about a conversation across all storage tiers"
|
||||
)
|
||||
async def get_conversation_stats(conversation_id: str):
|
||||
"""
|
||||
Get conversation statistics
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Statistics including turn counts and token usage across tiers
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
stats = await manager.get_conversation_stats(conversation_id)
|
||||
|
||||
return ConversationStatsResponse(**stats)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{conversation_id}/search",
|
||||
response_model=SearchResponse,
|
||||
summary="Search conversation semantically",
|
||||
description="Search for relevant turns within a conversation using semantic similarity"
|
||||
)
|
||||
async def search_conversation(
|
||||
conversation_id: str,
|
||||
search_request: SearchRequest
|
||||
):
|
||||
"""
|
||||
Semantic search within a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
search_request: Search query and parameters
|
||||
|
||||
Returns:
|
||||
Relevant conversation turns ranked by semantic similarity
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Perform semantic search
|
||||
results = await manager.search_conversations(
|
||||
query=search_request.query,
|
||||
conversation_id=conversation_id,
|
||||
limit=search_request.limit
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
search_results = [
|
||||
SearchResultResponse(
|
||||
conversation_id=result["conversation_id"],
|
||||
turn_number=result["turn_number"],
|
||||
role=result["role"],
|
||||
content=result["content"],
|
||||
timestamp=result["timestamp"],
|
||||
score=result["score"]
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
return SearchResponse(
|
||||
query=search_request.query,
|
||||
results=search_results,
|
||||
count=len(search_results)
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/search",
|
||||
response_model=SearchResponse,
|
||||
summary="Search all conversations",
|
||||
description="Search across all conversations using semantic similarity"
|
||||
)
|
||||
async def search_all_conversations(search_request: SearchRequest):
|
||||
"""
|
||||
Semantic search across all conversations
|
||||
|
||||
Args:
|
||||
search_request: Search query and parameters
|
||||
|
||||
Returns:
|
||||
Relevant turns from any conversation ranked by semantic similarity
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
# Perform semantic search across all conversations
|
||||
results = await manager.search_conversations(
|
||||
query=search_request.query,
|
||||
conversation_id=None, # Search all conversations
|
||||
limit=search_request.limit
|
||||
)
|
||||
|
||||
# Convert to response format
|
||||
search_results = [
|
||||
SearchResultResponse(
|
||||
conversation_id=result["conversation_id"],
|
||||
turn_number=result["turn_number"],
|
||||
role=result["role"],
|
||||
content=result["content"],
|
||||
timestamp=result["timestamp"],
|
||||
score=result["score"]
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
return SearchResponse(
|
||||
query=search_request.query,
|
||||
results=search_results,
|
||||
count=len(search_results)
|
||||
)
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/{conversation_id}",
|
||||
response_model=DeleteResponse,
|
||||
summary="Delete conversation",
|
||||
description="Delete a conversation from all storage tiers"
|
||||
)
|
||||
async def delete_conversation(
|
||||
conversation_id: str,
|
||||
clear_buffer: bool = Query(True, description="Clear from buffer (Tier 1)"),
|
||||
clear_qdrant: bool = Query(True, description="Clear from Qdrant (Tier 2/3)")
|
||||
):
|
||||
"""
|
||||
Delete a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
clear_buffer: Clear from Tier 1 buffer
|
||||
clear_qdrant: Clear from Tier 2/3 Qdrant
|
||||
|
||||
Returns:
|
||||
Deletion confirmation
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
try:
|
||||
await manager.clear_conversation(
|
||||
conversation_id,
|
||||
clear_buffer=clear_buffer,
|
||||
clear_qdrant=clear_qdrant
|
||||
)
|
||||
|
||||
return DeleteResponse(
|
||||
conversation_id=conversation_id,
|
||||
deleted=True,
|
||||
message=f"Conversation {conversation_id} deleted successfully"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error deleting conversation: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{conversation_id}/consolidate",
|
||||
summary="Consolidate conversation",
|
||||
description="Manually trigger consolidation from buffer to persistent storage"
|
||||
)
|
||||
async def consolidate_conversation(conversation_id: str):
|
||||
"""
|
||||
Manually consolidate a conversation
|
||||
|
||||
Moves all buffer turns to Qdrant persistent storage.
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Number of turns consolidated
|
||||
"""
|
||||
manager = get_memory_manager()
|
||||
|
||||
try:
|
||||
count = await manager.consolidate(conversation_id)
|
||||
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
"consolidated_turns": count,
|
||||
"message": f"Successfully consolidated {count} turns to persistent storage"
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error consolidating conversation: {str(e)}"
|
||||
)
|
||||
@@ -0,0 +1,34 @@
|
||||
"""
|
||||
OpenAI-compatible /v1/models endpoint
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter
|
||||
from .schemas import ModelsListResponse, ModelInfo
|
||||
from src.config import get_settings
|
||||
|
||||
router = APIRouter()
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
@router.get("/v1/models")
|
||||
async def list_models():
|
||||
"""List available models in OpenAI format."""
|
||||
|
||||
models = []
|
||||
|
||||
# Add OpenAI-style aliases
|
||||
for alias in settings.model_aliases.keys():
|
||||
models.append(ModelInfo(id=alias, owned_by="tatlock"))
|
||||
|
||||
# Add actual local models
|
||||
for model_list in [
|
||||
settings.get_lightweight_models(),
|
||||
settings.get_heavy_models(),
|
||||
settings.get_code_models()
|
||||
]:
|
||||
for model in model_list:
|
||||
# Avoid duplicates
|
||||
if model not in [m.id for m in models]:
|
||||
models.append(ModelInfo(id=model, owned_by="tatlock"))
|
||||
|
||||
return ModelsListResponse(data=models)
|
||||
@@ -0,0 +1,142 @@
|
||||
"""
|
||||
OpenAI-compatible API schemas for /v1/* endpoints
|
||||
"""
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Literal
|
||||
from enum import Enum
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Request Schemas
|
||||
# ============================================================================
|
||||
|
||||
class MessageRole(str, Enum):
|
||||
"""Valid message roles."""
|
||||
SYSTEM = "system"
|
||||
USER = "user"
|
||||
ASSISTANT = "assistant"
|
||||
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
"""A single message in the conversation."""
|
||||
role: MessageRole
|
||||
content: str
|
||||
name: Optional[str] = None
|
||||
|
||||
|
||||
class ChatCompletionRequest(BaseModel):
|
||||
"""OpenAI-compatible chat completion request."""
|
||||
|
||||
model: str = Field(..., description="Model to use")
|
||||
messages: List[ChatMessage] = Field(..., min_length=1)
|
||||
stream: bool = Field(default=False, description="Enable streaming")
|
||||
|
||||
# Memory system (Phase 2)
|
||||
conversation_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Conversation ID for memory tracking (auto-generated if not provided)"
|
||||
)
|
||||
store_in_memory: bool = Field(
|
||||
default=True,
|
||||
description="Store conversation turns in memory system"
|
||||
)
|
||||
|
||||
# Optional parameters
|
||||
temperature: Optional[float] = Field(default=0.7, ge=0, le=2)
|
||||
top_p: Optional[float] = Field(default=1.0, ge=0, le=1)
|
||||
max_tokens: Optional[int] = Field(default=None, ge=1)
|
||||
frequency_penalty: Optional[float] = Field(default=0.0, ge=-2, le=2)
|
||||
presence_penalty: Optional[float] = Field(default=0.0, ge=-2, le=2)
|
||||
stop: Optional[List[str]] = None
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello!"}
|
||||
],
|
||||
"stream": False,
|
||||
"temperature": 0.7
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Response Schemas
|
||||
# ============================================================================
|
||||
|
||||
class ChatMessageResponse(BaseModel):
|
||||
"""Response message."""
|
||||
role: str = "assistant"
|
||||
content: str
|
||||
|
||||
|
||||
class ChatCompletionChoice(BaseModel):
|
||||
"""A single completion choice."""
|
||||
index: int = 0
|
||||
message: ChatMessageResponse
|
||||
finish_reason: str = "stop"
|
||||
|
||||
|
||||
class UsageInfo(BaseModel):
|
||||
"""Token usage information."""
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
|
||||
|
||||
class ChatCompletionResponse(BaseModel):
|
||||
"""OpenAI-compatible chat completion response (non-streaming)."""
|
||||
|
||||
id: str
|
||||
object: str = "chat.completion"
|
||||
created: int
|
||||
model: str
|
||||
choices: List[ChatCompletionChoice]
|
||||
usage: UsageInfo
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Streaming Response Schemas
|
||||
# ============================================================================
|
||||
|
||||
class DeltaMessage(BaseModel):
|
||||
"""Delta message for streaming."""
|
||||
role: Optional[str] = None
|
||||
content: Optional[str] = None
|
||||
|
||||
|
||||
class ChatCompletionStreamChoice(BaseModel):
|
||||
"""Streaming choice."""
|
||||
index: int = 0
|
||||
delta: DeltaMessage
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
|
||||
class ChatCompletionStreamResponse(BaseModel):
|
||||
"""OpenAI-compatible streaming chunk."""
|
||||
|
||||
id: str
|
||||
object: str = "chat.completion.chunk"
|
||||
created: int
|
||||
model: str
|
||||
choices: List[ChatCompletionStreamChoice]
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Models Endpoint
|
||||
# ============================================================================
|
||||
|
||||
class ModelInfo(BaseModel):
|
||||
"""Model information."""
|
||||
id: str
|
||||
object: str = "model"
|
||||
created: int = 0
|
||||
owned_by: str = "local"
|
||||
|
||||
|
||||
class ModelsListResponse(BaseModel):
|
||||
"""List of available models."""
|
||||
object: str = "list"
|
||||
data: List[ModelInfo]
|
||||
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
Base Pydantic models for consistent schema behavior
|
||||
"""
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
|
||||
class BaseSchema(BaseModel):
|
||||
"""
|
||||
Base Pydantic model with standardized configuration
|
||||
|
||||
All schemas should inherit from this to ensure consistent behavior:
|
||||
- Consistent datetime serialization
|
||||
- Strict validation by default
|
||||
- JSON schema generation
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
# Strict type validation
|
||||
strict=False,
|
||||
|
||||
# Allow population by field name
|
||||
populate_by_name=True,
|
||||
|
||||
# Use enum values in JSON
|
||||
use_enum_values=True,
|
||||
|
||||
# Validate assignments after initialization
|
||||
validate_assignment=True,
|
||||
|
||||
# Serialize datetime to ISO format
|
||||
json_encoders={
|
||||
datetime: lambda v: v.isoformat() if v else None
|
||||
}
|
||||
)
|
||||
|
||||
def dict_without_none(self) -> dict[str, Any]:
|
||||
"""
|
||||
Return model as dict, excluding None values
|
||||
|
||||
Returns:
|
||||
Dictionary with None values filtered out
|
||||
"""
|
||||
return {k: v for k, v in self.model_dump().items() if v is not None}
|
||||
@@ -0,0 +1,5 @@
|
||||
"""
|
||||
API Clients package for Core-API
|
||||
|
||||
Provides HTTP/WebSocket clients for external infrastructure services.
|
||||
"""
|
||||
@@ -0,0 +1,271 @@
|
||||
"""
|
||||
Nginx Proxy Manager API Client
|
||||
|
||||
Provides interface to NPM REST API for proxy host and SSL certificate management.
|
||||
"""
|
||||
import httpx
|
||||
from typing import Optional, Dict, List, Any
|
||||
from datetime import datetime, timedelta
|
||||
from src.logging_config import get_logger
|
||||
from src.config import get_settings
|
||||
|
||||
logger = get_logger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class NPMClient:
|
||||
"""
|
||||
HTTP client for Nginx Proxy Manager API
|
||||
|
||||
Uses JWT Bearer token authentication with automatic token refresh.
|
||||
Tokens expire after ~24 hours.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: Optional[str] = None,
|
||||
email: Optional[str] = None,
|
||||
password: Optional[str] = None,
|
||||
timeout: int = 30
|
||||
):
|
||||
"""
|
||||
Initialize NPM client
|
||||
|
||||
Args:
|
||||
base_url: NPM base URL (default from settings)
|
||||
email: NPM admin email (default from settings)
|
||||
password: NPM admin password (default from settings)
|
||||
timeout: Request timeout in seconds
|
||||
"""
|
||||
self.base_url = (base_url or settings.npm_url).rstrip("/")
|
||||
self.email = email or settings.npm_email
|
||||
self.password = password or settings.npm_password
|
||||
self.timeout = timeout
|
||||
|
||||
self._token: Optional[str] = None
|
||||
self._token_expires: Optional[datetime] = None
|
||||
|
||||
if not self.email or not self.password:
|
||||
logger.warning("NPM credentials not configured")
|
||||
|
||||
async def _ensure_token(self):
|
||||
"""Ensure we have a valid token, refresh if needed"""
|
||||
if self._token and self._token_expires:
|
||||
# If token expires in less than 1 hour, refresh it
|
||||
if datetime.now() + timedelta(hours=1) < self._token_expires:
|
||||
return
|
||||
|
||||
# Get new token
|
||||
await self._refresh_token()
|
||||
|
||||
async def _refresh_token(self):
|
||||
"""Get a new authentication token"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/tokens",
|
||||
json={
|
||||
"identity": self.email,
|
||||
"secret": self.password
|
||||
}
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
self._token = data.get("token")
|
||||
# Assume 23-hour expiration to be safe
|
||||
self._token_expires = datetime.now() + timedelta(hours=23)
|
||||
|
||||
logger.info("NPM token refreshed successfully")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to refresh NPM token: {e}")
|
||||
raise
|
||||
|
||||
def _get_headers(self) -> Dict[str, str]:
|
||||
"""Get request headers with authentication"""
|
||||
if not self._token:
|
||||
raise RuntimeError("No NPM token available. Call _ensure_token() first.")
|
||||
|
||||
return {
|
||||
"Authorization": f"Bearer {self._token}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""
|
||||
Check if NPM API is accessible
|
||||
|
||||
Returns:
|
||||
True if accessible, False otherwise
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.get(f"{self.base_url}/api")
|
||||
return response.status_code == 200
|
||||
except Exception as e:
|
||||
logger.error(f"NPM health check failed: {e}")
|
||||
return False
|
||||
|
||||
async def get_proxy_hosts(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all proxy hosts
|
||||
|
||||
Returns:
|
||||
List of proxy host configurations
|
||||
"""
|
||||
await self._ensure_token()
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/nginx/proxy-hosts",
|
||||
headers=self._get_headers()
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def get_proxy_host(self, host_id: int) -> Dict[str, Any]:
|
||||
"""
|
||||
Get details of a specific proxy host
|
||||
|
||||
Args:
|
||||
host_id: Proxy host identifier
|
||||
|
||||
Returns:
|
||||
Proxy host configuration
|
||||
"""
|
||||
await self._ensure_token()
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/nginx/proxy-hosts/{host_id}",
|
||||
headers=self._get_headers()
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def create_proxy_host(
|
||||
self,
|
||||
domain_names: List[str],
|
||||
forward_host: str,
|
||||
forward_port: int,
|
||||
forward_scheme: str = "http",
|
||||
certificate_id: int = 0,
|
||||
ssl_forced: bool = False,
|
||||
block_exploits: bool = True,
|
||||
caching_enabled: bool = True,
|
||||
websocket_upgrade: bool = True,
|
||||
http2_support: bool = True,
|
||||
hsts_enabled: bool = True,
|
||||
advanced_config: str = ""
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Create a new proxy host
|
||||
|
||||
Args:
|
||||
domain_names: List of domain names for this proxy
|
||||
forward_host: Target host to proxy to
|
||||
forward_port: Target port to proxy to
|
||||
forward_scheme: http or https
|
||||
certificate_id: SSL certificate ID (0 for none)
|
||||
ssl_forced: Force HTTPS redirect
|
||||
block_exploits: Enable exploit blocking
|
||||
caching_enabled: Enable response caching
|
||||
websocket_upgrade: Allow WebSocket upgrades
|
||||
http2_support: Enable HTTP/2
|
||||
hsts_enabled: Enable HSTS headers
|
||||
advanced_config: Custom nginx configuration
|
||||
|
||||
Returns:
|
||||
Created proxy host details
|
||||
"""
|
||||
await self._ensure_token()
|
||||
|
||||
payload = {
|
||||
"domain_names": domain_names,
|
||||
"forward_scheme": forward_scheme,
|
||||
"forward_host": forward_host,
|
||||
"forward_port": forward_port,
|
||||
"certificate_id": certificate_id,
|
||||
"ssl_forced": ssl_forced,
|
||||
"block_exploits": block_exploits,
|
||||
"caching_enabled": caching_enabled,
|
||||
"allow_websocket_upgrade": websocket_upgrade,
|
||||
"http2_support": http2_support,
|
||||
"hsts_enabled": hsts_enabled,
|
||||
"hsts_subdomains": False,
|
||||
"advanced_config": advanced_config,
|
||||
"access_list_id": 0,
|
||||
"meta": {}
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/nginx/proxy-hosts",
|
||||
headers=self._get_headers(),
|
||||
json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def get_certificates(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all SSL certificates
|
||||
|
||||
Returns:
|
||||
List of certificate details
|
||||
"""
|
||||
await self._ensure_token()
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/nginx/certificates",
|
||||
headers=self._get_headers()
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def create_certificate(
|
||||
self,
|
||||
domain_names: List[str],
|
||||
provider: str = "letsencrypt"
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Request a new SSL certificate from Let's Encrypt
|
||||
|
||||
Args:
|
||||
domain_names: List of domains for the certificate
|
||||
provider: Certificate provider (default: letsencrypt)
|
||||
|
||||
Returns:
|
||||
Certificate details
|
||||
"""
|
||||
await self._ensure_token()
|
||||
|
||||
payload = {
|
||||
"provider": provider,
|
||||
"domain_names": domain_names,
|
||||
"meta": {
|
||||
"dns_challenge": False
|
||||
}
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/nginx/certificates",
|
||||
headers=self._get_headers(),
|
||||
json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_npm_client: Optional[NPMClient] = None
|
||||
|
||||
|
||||
def get_npm_client() -> NPMClient:
|
||||
"""Get singleton NPM client instance"""
|
||||
global _npm_client
|
||||
if _npm_client is None:
|
||||
_npm_client = NPMClient()
|
||||
return _npm_client
|
||||
@@ -0,0 +1,221 @@
|
||||
"""
|
||||
Portainer API Client
|
||||
|
||||
Provides interface to Portainer REST API for stack and container management.
|
||||
"""
|
||||
import httpx
|
||||
from typing import Optional, Dict, List, Any
|
||||
from src.logging_config import get_logger
|
||||
from src.config import get_settings
|
||||
|
||||
logger = get_logger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class PortainerClient:
|
||||
"""
|
||||
HTTP client for Portainer API
|
||||
|
||||
Uses access token authentication (X-API-Key header)
|
||||
for long-lived API access without session management.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
timeout: int = 30
|
||||
):
|
||||
"""
|
||||
Initialize Portainer client
|
||||
|
||||
Args:
|
||||
base_url: Portainer base URL (default from settings)
|
||||
api_key: Portainer API access token (default from settings)
|
||||
timeout: Request timeout in seconds
|
||||
"""
|
||||
self.base_url = (base_url or settings.portainer_url).rstrip("/")
|
||||
self.api_key = api_key or settings.portainer_api_key
|
||||
self.timeout = timeout
|
||||
|
||||
if not self.api_key:
|
||||
logger.warning("Portainer API key not configured")
|
||||
|
||||
def _get_headers(self) -> Dict[str, str]:
|
||||
"""Get request headers with authentication"""
|
||||
return {
|
||||
"X-API-Key": self.api_key,
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""
|
||||
Check if Portainer API is accessible
|
||||
|
||||
Returns:
|
||||
True if accessible, False otherwise
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.get(f"{self.base_url}/api/status")
|
||||
return response.status_code == 200
|
||||
except Exception as e:
|
||||
logger.error(f"Portainer health check failed: {e}")
|
||||
return False
|
||||
|
||||
async def get_endpoints(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all Portainer endpoints (Docker environments)
|
||||
|
||||
Returns:
|
||||
List of endpoint configurations
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/endpoints",
|
||||
headers=self._get_headers()
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def get_stacks(self, endpoint_id: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all stacks
|
||||
|
||||
Args:
|
||||
endpoint_id: Filter by specific endpoint (optional)
|
||||
|
||||
Returns:
|
||||
List of stack configurations
|
||||
"""
|
||||
params = {}
|
||||
if endpoint_id:
|
||||
params["endpointId"] = endpoint_id
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/stacks",
|
||||
headers=self._get_headers(),
|
||||
params=params
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def get_stack(self, stack_id: int) -> Dict[str, Any]:
|
||||
"""
|
||||
Get details of a specific stack
|
||||
|
||||
Args:
|
||||
stack_id: Stack identifier
|
||||
|
||||
Returns:
|
||||
Stack configuration details
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/stacks/{stack_id}",
|
||||
headers=self._get_headers()
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def create_stack(
|
||||
self,
|
||||
name: str,
|
||||
stack_file_content: str,
|
||||
endpoint_id: int
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Create a new stack from compose file content
|
||||
|
||||
Args:
|
||||
name: Stack name
|
||||
stack_file_content: Docker Compose YAML content
|
||||
endpoint_id: Portainer endpoint to deploy to
|
||||
|
||||
Returns:
|
||||
Created stack details
|
||||
"""
|
||||
payload = {
|
||||
"name": name,
|
||||
"stackFileContent": stack_file_content
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/stacks/create/standalone/string",
|
||||
headers=self._get_headers(),
|
||||
params={"endpointId": endpoint_id},
|
||||
json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def update_stack(
|
||||
self,
|
||||
stack_id: int,
|
||||
stack_file_content: str,
|
||||
endpoint_id: int,
|
||||
prune: bool = False,
|
||||
pull_image: bool = False
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Update an existing stack
|
||||
|
||||
Args:
|
||||
stack_id: Stack identifier
|
||||
stack_file_content: New Docker Compose YAML content
|
||||
endpoint_id: Portainer endpoint
|
||||
prune: Remove services no longer defined
|
||||
pull_image: Pull latest images before deployment
|
||||
|
||||
Returns:
|
||||
Updated stack details
|
||||
"""
|
||||
payload = {
|
||||
"stackFileContent": stack_file_content,
|
||||
"prune": prune,
|
||||
"pullImage": pull_image
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.put(
|
||||
f"{self.base_url}/api/stacks/{stack_id}",
|
||||
headers=self._get_headers(),
|
||||
params={"endpointId": endpoint_id},
|
||||
json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def delete_stack(self, stack_id: int, endpoint_id: int) -> bool:
|
||||
"""
|
||||
Delete a stack
|
||||
|
||||
Args:
|
||||
stack_id: Stack identifier
|
||||
endpoint_id: Portainer endpoint
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
response = await client.delete(
|
||||
f"{self.base_url}/api/stacks/{stack_id}",
|
||||
headers=self._get_headers(),
|
||||
params={"endpointId": endpoint_id}
|
||||
)
|
||||
response.raise_for_status()
|
||||
return True
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_portainer_client: Optional[PortainerClient] = None
|
||||
|
||||
|
||||
def get_portainer_client() -> PortainerClient:
|
||||
"""Get singleton Portainer client instance"""
|
||||
global _portainer_client
|
||||
if _portainer_client is None:
|
||||
_portainer_client = PortainerClient()
|
||||
return _portainer_client
|
||||
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Global configuration for Core Code API
|
||||
"""
|
||||
from pydantic_settings import BaseSettings
|
||||
from functools import lru_cache
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""Global application settings"""
|
||||
|
||||
# Application
|
||||
app_name: str = "Core Code API"
|
||||
app_version: str = "1.0.0"
|
||||
debug: bool = False
|
||||
|
||||
# Server
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8083
|
||||
|
||||
# CORS
|
||||
cors_origins: list[str] = ["*"]
|
||||
cors_credentials: bool = True
|
||||
cors_methods: list[str] = ["*"]
|
||||
cors_headers: list[str] = ["*"]
|
||||
|
||||
# Logging
|
||||
log_level: str = "INFO"
|
||||
|
||||
# Ollama Configuration (for AI orchestration)
|
||||
ollama_base_url: str = "http://ollama:11434"
|
||||
ollama_timeout: int = 300 # 5 minutes
|
||||
|
||||
# Model Configuration
|
||||
default_model: str = "gemma:7b"
|
||||
lightweight_models: str = "gemma:2b,gemma:7b"
|
||||
heavy_models: str = "mistral:7b,gemma2:9b,mixtral:8x7b"
|
||||
code_models: str = "codestral:latest,codegemma:latest"
|
||||
|
||||
# Model Aliases (OpenAI → Local)
|
||||
alias_gpt35: str = "gemma:7b"
|
||||
alias_gpt4: str = "mistral:7b"
|
||||
alias_gpt4_turbo: str = "mixtral:8x7b"
|
||||
alias_gpt4_code: str = "codestral:latest"
|
||||
|
||||
# Memory Configuration
|
||||
memory_tier1_max_turns: int = 10
|
||||
memory_consolidation_threshold: int = 10
|
||||
|
||||
# Qdrant Configuration
|
||||
qdrant_host: str = "qdrant"
|
||||
qdrant_port: int = 6333
|
||||
qdrant_collection_conversations: str = "core_api_conversations"
|
||||
qdrant_collection_documents: str = "core_api_documents"
|
||||
qdrant_collection_user_facts: str = "core_api_user_facts"
|
||||
|
||||
# Embeddings
|
||||
embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
|
||||
embedding_dimension: int = 384
|
||||
embedding_batch_size: int = 32
|
||||
|
||||
# Infrastructure Management
|
||||
portainer_url: str = "http://localhost:8001"
|
||||
portainer_api_key: str = ""
|
||||
|
||||
npm_url: str = "http://localhost:81"
|
||||
npm_email: str = ""
|
||||
npm_password: str = ""
|
||||
|
||||
kuma_url: str = "http://localhost:3001"
|
||||
kuma_username: str = ""
|
||||
kuma_password: str = ""
|
||||
|
||||
@property
|
||||
def model_aliases(self) -> dict:
|
||||
"""Computed property for model aliases"""
|
||||
return {
|
||||
"gpt-3.5-turbo": self.alias_gpt35,
|
||||
"gpt-4": self.alias_gpt4,
|
||||
"gpt-4-turbo": self.alias_gpt4_turbo,
|
||||
"gpt-4-code": self.alias_gpt4_code,
|
||||
}
|
||||
|
||||
def get_lightweight_models(self) -> list[str]:
|
||||
"""Parse comma-separated lightweight models"""
|
||||
return [m.strip().strip('"').strip("'") for m in self.lightweight_models.split(",") if m.strip()]
|
||||
|
||||
def get_heavy_models(self) -> list[str]:
|
||||
"""Parse comma-separated heavy models"""
|
||||
return [m.strip().strip('"').strip("'") for m in self.heavy_models.split(",") if m.strip()]
|
||||
|
||||
def get_code_models(self) -> list[str]:
|
||||
"""Parse comma-separated code models"""
|
||||
return [m.strip().strip('"').strip("'") for m in self.code_models.split(",") if m.strip()]
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
case_sensitive = False
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def get_settings() -> Settings:
|
||||
"""Cached settings instance"""
|
||||
return Settings()
|
||||
@@ -0,0 +1,5 @@
|
||||
"""
|
||||
Controllers package for Core-API
|
||||
|
||||
Provides controller-based routing architecture for better code organization.
|
||||
"""
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
Base controller class for Core-API
|
||||
|
||||
Provides common functionality for all controllers.
|
||||
"""
|
||||
from fastapi import APIRouter
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class BaseController(ABC):
|
||||
"""
|
||||
Base controller class with common functionality
|
||||
|
||||
All controllers should inherit from this class and implement
|
||||
the create_router() method to define their endpoints.
|
||||
"""
|
||||
|
||||
def __init__(self, prefix: str, tags: list[str]):
|
||||
"""
|
||||
Initialize base controller
|
||||
|
||||
Args:
|
||||
prefix: URL prefix for this controller's routes
|
||||
tags: OpenAPI tags for documentation grouping
|
||||
"""
|
||||
self.prefix = prefix
|
||||
self.tags = tags
|
||||
self._router = None
|
||||
|
||||
@abstractmethod
|
||||
def create_router(self) -> APIRouter:
|
||||
"""
|
||||
Create and configure the FastAPI router for this controller
|
||||
|
||||
Returns:
|
||||
Configured APIRouter instance with all endpoints
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
def router(self) -> APIRouter:
|
||||
"""
|
||||
Get the router instance, creating it if needed
|
||||
|
||||
Returns:
|
||||
APIRouter instance
|
||||
"""
|
||||
if self._router is None:
|
||||
self._router = self.create_router()
|
||||
return self._router
|
||||
@@ -0,0 +1,270 @@
|
||||
"""
|
||||
Infrastructure Management Controller
|
||||
|
||||
Provides API endpoints for automated infrastructure management,
|
||||
including service deployment, configuration, and monitoring setup.
|
||||
"""
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from typing import List, Dict, Any, Optional
|
||||
from pydantic import BaseModel
|
||||
|
||||
from src.controllers.base import BaseController
|
||||
from src.clients.portainer_client import get_portainer_client
|
||||
from src.clients.npm_client import get_npm_client
|
||||
from src.logging_config import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# Response models
|
||||
class ServiceInfo(BaseModel):
|
||||
"""Information about a deployed service"""
|
||||
name: str
|
||||
stack_id: Optional[int]
|
||||
status: str
|
||||
endpoint_id: Optional[int]
|
||||
ports: List[int] = []
|
||||
domains: List[str] = []
|
||||
|
||||
|
||||
class PortInfo(BaseModel):
|
||||
"""Information about an allocated port"""
|
||||
port: int
|
||||
service: str
|
||||
protocol: str = "tcp"
|
||||
description: str = ""
|
||||
|
||||
|
||||
class DomainInfo(BaseModel):
|
||||
"""Information about a configured domain"""
|
||||
domain: str
|
||||
service: str
|
||||
proxy_host_id: Optional[int]
|
||||
ssl_enabled: bool = False
|
||||
certificate_id: Optional[int]
|
||||
|
||||
|
||||
class InfrastructureHealth(BaseModel):
|
||||
"""Overall infrastructure health status"""
|
||||
portainer_connected: bool
|
||||
npm_connected: bool
|
||||
total_stacks: int
|
||||
total_proxy_hosts: int
|
||||
|
||||
|
||||
class InfrastructureController(BaseController):
|
||||
"""
|
||||
Controller for infrastructure management operations
|
||||
|
||||
Provides endpoints for:
|
||||
- Service discovery and listing
|
||||
- Port allocation management
|
||||
- Domain/proxy configuration
|
||||
- Automated service deployment
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(prefix="/infrastructure", tags=["Infrastructure"])
|
||||
|
||||
def create_router(self) -> APIRouter:
|
||||
"""Create and configure the router"""
|
||||
router = APIRouter(prefix=self.prefix, tags=self.tags)
|
||||
|
||||
@router.get(
|
||||
"/health",
|
||||
response_model=InfrastructureHealth,
|
||||
summary="Infrastructure health check"
|
||||
)
|
||||
async def get_infrastructure_health():
|
||||
"""
|
||||
Check health of all infrastructure services
|
||||
|
||||
Returns status of Portainer, NPM, and summary statistics.
|
||||
"""
|
||||
portainer = get_portainer_client()
|
||||
npm = get_npm_client()
|
||||
|
||||
portainer_healthy = await portainer.health_check()
|
||||
npm_healthy = await npm.health_check()
|
||||
|
||||
total_stacks = 0
|
||||
total_proxy_hosts = 0
|
||||
|
||||
if portainer_healthy:
|
||||
try:
|
||||
stacks = await portainer.get_stacks()
|
||||
total_stacks = len(stacks)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get stacks count: {e}")
|
||||
|
||||
if npm_healthy:
|
||||
try:
|
||||
proxy_hosts = await npm.get_proxy_hosts()
|
||||
total_proxy_hosts = len(proxy_hosts)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get proxy hosts count: {e}")
|
||||
|
||||
return InfrastructureHealth(
|
||||
portainer_connected=portainer_healthy,
|
||||
npm_connected=npm_healthy,
|
||||
total_stacks=total_stacks,
|
||||
total_proxy_hosts=total_proxy_hosts
|
||||
)
|
||||
|
||||
@router.get(
|
||||
"/services",
|
||||
response_model=List[ServiceInfo],
|
||||
summary="List all deployed services"
|
||||
)
|
||||
async def list_services():
|
||||
"""
|
||||
List all deployed services from Portainer stacks
|
||||
|
||||
Returns comprehensive service information including:
|
||||
- Stack/service name
|
||||
- Status
|
||||
- Exposed ports
|
||||
- Configured domains
|
||||
"""
|
||||
portainer = get_portainer_client()
|
||||
npm = get_npm_client()
|
||||
|
||||
try:
|
||||
stacks = await portainer.get_stacks()
|
||||
proxy_hosts = await npm.get_proxy_hosts()
|
||||
|
||||
# Build domain mapping (domain -> service name)
|
||||
domain_map = {}
|
||||
for proxy in proxy_hosts:
|
||||
for domain in proxy.get("domain_names", []):
|
||||
# Try to extract service name from forward_host
|
||||
forward_host = proxy.get("forward_host", "")
|
||||
domain_map[domain] = forward_host
|
||||
|
||||
services = []
|
||||
for stack in stacks:
|
||||
# Find domains for this stack
|
||||
stack_name = stack.get("Name", "")
|
||||
domains = [
|
||||
domain for domain, host in domain_map.items()
|
||||
if stack_name in host or host in stack_name
|
||||
]
|
||||
|
||||
service_info = ServiceInfo(
|
||||
name=stack_name,
|
||||
stack_id=stack.get("Id"),
|
||||
status=stack.get("Status", "unknown"),
|
||||
endpoint_id=stack.get("EndpointId"),
|
||||
ports=[], # TODO: Extract from stack file
|
||||
domains=domains
|
||||
)
|
||||
services.append(service_info)
|
||||
|
||||
return services
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list services: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@router.get(
|
||||
"/services/{name}",
|
||||
response_model=ServiceInfo,
|
||||
summary="Get service details"
|
||||
)
|
||||
async def get_service(name: str):
|
||||
"""
|
||||
Get detailed information about a specific service
|
||||
|
||||
Args:
|
||||
name: Service/stack name
|
||||
|
||||
Returns:
|
||||
Service details including status and configuration
|
||||
"""
|
||||
portainer = get_portainer_client()
|
||||
|
||||
try:
|
||||
stacks = await portainer.get_stacks()
|
||||
|
||||
# Find stack by name (case-insensitive)
|
||||
stack = next(
|
||||
(s for s in stacks if s.get("Name", "").lower() == name.lower()),
|
||||
None
|
||||
)
|
||||
|
||||
if not stack:
|
||||
raise HTTPException(status_code=404, detail=f"Service '{name}' not found")
|
||||
|
||||
return ServiceInfo(
|
||||
name=stack.get("Name", ""),
|
||||
stack_id=stack.get("Id"),
|
||||
status=stack.get("Status", "unknown"),
|
||||
endpoint_id=stack.get("EndpointId"),
|
||||
ports=[],
|
||||
domains=[]
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get service '{name}': {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@router.get(
|
||||
"/ports",
|
||||
response_model=List[PortInfo],
|
||||
summary="List allocated ports"
|
||||
)
|
||||
async def list_ports():
|
||||
"""
|
||||
List all currently allocated ports
|
||||
|
||||
Scans services and proxy configurations to build
|
||||
a comprehensive port allocation map.
|
||||
"""
|
||||
# TODO: Implement port scanning from containers and proxy configs
|
||||
# For now, return a placeholder
|
||||
return []
|
||||
|
||||
@router.get(
|
||||
"/domains",
|
||||
response_model=List[DomainInfo],
|
||||
summary="List configured domains"
|
||||
)
|
||||
async def list_domains():
|
||||
"""
|
||||
List all configured domain names
|
||||
|
||||
Returns domain-to-service mappings with SSL status.
|
||||
"""
|
||||
npm = get_npm_client()
|
||||
|
||||
try:
|
||||
proxy_hosts = await npm.get_proxy_hosts()
|
||||
|
||||
domains = []
|
||||
for proxy in proxy_hosts:
|
||||
service_name = proxy.get("forward_host", "localhost")
|
||||
certificate_id = proxy.get("certificate_id", 0)
|
||||
|
||||
for domain in proxy.get("domain_names", []):
|
||||
domain_info = DomainInfo(
|
||||
domain=domain,
|
||||
service=service_name,
|
||||
proxy_host_id=proxy.get("id"),
|
||||
ssl_enabled=certificate_id > 0,
|
||||
certificate_id=certificate_id if certificate_id > 0 else None
|
||||
)
|
||||
domains.append(domain_info)
|
||||
|
||||
return domains
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list domains: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return router
|
||||
|
||||
|
||||
# Create controller instance
|
||||
infrastructure_controller = InfrastructureController()
|
||||
@@ -0,0 +1,24 @@
|
||||
"""
|
||||
Infrastructure Credentials Template
|
||||
|
||||
INSTRUCTIONS:
|
||||
1. Copy this file to credentials.py
|
||||
2. Fill in your actual credentials
|
||||
3. DO NOT commit credentials.py to version control (it's in .gitignore)
|
||||
|
||||
This file should be committed to the repository as a template.
|
||||
"""
|
||||
|
||||
# Portainer Configuration
|
||||
PORTAINER_URL = "http://localhost:8001"
|
||||
PORTAINER_API_KEY = "ptr_your_api_token_here" # Create in Portainer UI: User menu → My account → Access tokens
|
||||
|
||||
# Nginx Proxy Manager Configuration
|
||||
NPM_URL = "http://localhost:81"
|
||||
NPM_EMAIL = "admin@example.com"
|
||||
NPM_PASSWORD = "your_password_here"
|
||||
|
||||
# Uptime Kuma Configuration
|
||||
KUMA_URL = "http://localhost:3001"
|
||||
KUMA_USERNAME = "admin"
|
||||
KUMA_PASSWORD = "your_password_here"
|
||||
@@ -0,0 +1,49 @@
|
||||
"""
|
||||
Logging configuration for Core Code API
|
||||
"""
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def setup_logging(log_level: str = "INFO") -> None:
|
||||
"""
|
||||
Configure logging for the application
|
||||
|
||||
Args:
|
||||
log_level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
|
||||
"""
|
||||
# Create logs directory if it doesn't exist
|
||||
log_dir = Path("logs")
|
||||
log_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Configure root logger
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, log_level.upper()),
|
||||
format="%(asctime)s | %(levelname)-8s | %(name)s:%(funcName)s:%(lineno)d | %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
handlers=[
|
||||
# Console handler
|
||||
logging.StreamHandler(sys.stdout),
|
||||
# File handler
|
||||
logging.FileHandler(log_dir / "app.log", encoding="utf-8")
|
||||
]
|
||||
)
|
||||
|
||||
# Set specific log levels for third-party libraries
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
logging.getLogger("httpcore").setLevel(logging.WARNING)
|
||||
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
def get_logger(name: str) -> logging.Logger:
|
||||
"""
|
||||
Get a logger instance
|
||||
|
||||
Args:
|
||||
name: Logger name (typically __name__)
|
||||
|
||||
Returns:
|
||||
Configured logger instance
|
||||
"""
|
||||
return logging.getLogger(name)
|
||||
@@ -0,0 +1,200 @@
|
||||
"""
|
||||
Main FastAPI application for Core Code API
|
||||
"""
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import JSONResponse
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from src.config import get_settings
|
||||
from src.logging_config import setup_logging, get_logger
|
||||
from src.web_scraper import router as web_scraper_router
|
||||
from src.api.v1.chat import router as chat_router
|
||||
from src.api.v1.models import router as models_router
|
||||
from src.api.v1.conversations import router as conversations_router
|
||||
from src.models.ollama_client import get_ollama_client, close_ollama_client
|
||||
|
||||
# Initialize settings
|
||||
settings = get_settings()
|
||||
|
||||
# Setup logging
|
||||
setup_logging(settings.log_level)
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""
|
||||
Application lifespan manager for startup/shutdown events
|
||||
|
||||
Args:
|
||||
app: FastAPI application instance
|
||||
"""
|
||||
# Startup
|
||||
logger.info("=" * 60)
|
||||
logger.info(f"Starting {settings.app_name} v{settings.app_version}")
|
||||
logger.info(f"Debug mode: {settings.debug}")
|
||||
logger.info(f"Log level: {settings.log_level}")
|
||||
logger.info(f"Ollama URL: {settings.ollama_base_url}")
|
||||
logger.info("=" * 60)
|
||||
|
||||
# Check Ollama connectivity
|
||||
ollama_client = get_ollama_client()
|
||||
ollama_healthy = await ollama_client.health_check()
|
||||
if ollama_healthy:
|
||||
logger.info("✓ Ollama connection successful")
|
||||
else:
|
||||
logger.warning("✗ Ollama connection failed - AI features may not work")
|
||||
|
||||
yield
|
||||
|
||||
# Shutdown
|
||||
logger.info("Shutting down application")
|
||||
await close_ollama_client()
|
||||
|
||||
|
||||
# Create FastAPI application
|
||||
app = FastAPI(
|
||||
title=settings.app_name,
|
||||
version=settings.app_version,
|
||||
description="""
|
||||
Core Code API provides OpenAPI-compatible functions and AI orchestration for Open WebUI.
|
||||
|
||||
## Features
|
||||
|
||||
### OpenAI-Compatible API (v1)
|
||||
- `/v1/chat/completions` - Chat completions with streaming support
|
||||
- `/v1/models` - List available models
|
||||
Compatible with OpenAI client libraries and Open WebUI.
|
||||
|
||||
### Conversation Memory (Phase 2)
|
||||
- `/v1/conversations/{id}` - Get conversation history
|
||||
- `/v1/conversations/{id}/search` - Semantic search within conversation
|
||||
- `/v1/conversations/search` - Search across all conversations
|
||||
- `/v1/conversations/{id}/stats` - Get conversation statistics
|
||||
- `/v1/conversations/{id}/consolidate` - Manual consolidation
|
||||
- `DELETE /v1/conversations/{id}` - Delete conversation
|
||||
|
||||
Multi-tier memory system:
|
||||
- **Tier 1**: Fast in-memory buffer (last 10 turns)
|
||||
- **Tier 2/3**: Unified Qdrant storage (persistent + semantic search)
|
||||
|
||||
### Web Scraper
|
||||
Intelligent web scraping with main content extraction.
|
||||
Perfect for extracting articles, documentation, and blog posts for LLM consumption.
|
||||
|
||||
## Integration
|
||||
|
||||
This API is designed to integrate with:
|
||||
- **Open WebUI**: Direct OpenAI API compatibility
|
||||
- **Open WebUI Functions**: Import via OpenAPI spec
|
||||
- **Open WebUI Pipelines**: Use as data source
|
||||
- **LangChain**: Compatible with standard HTTP tools
|
||||
|
||||
## Documentation
|
||||
|
||||
- **OpenAPI Spec**: `/openapi.json`
|
||||
- **Swagger UI**: `/docs`
|
||||
- **ReDoc**: `/redoc`
|
||||
""",
|
||||
docs_url="/docs",
|
||||
redoc_url="/redoc",
|
||||
openapi_url="/openapi.json",
|
||||
lifespan=lifespan,
|
||||
debug=settings.debug
|
||||
)
|
||||
|
||||
# Add CORS middleware
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.cors_origins,
|
||||
allow_credentials=settings.cors_credentials,
|
||||
allow_methods=settings.cors_methods,
|
||||
allow_headers=settings.cors_headers,
|
||||
)
|
||||
|
||||
|
||||
# Root endpoint
|
||||
@app.get(
|
||||
"/",
|
||||
tags=["Health"],
|
||||
summary="Service information",
|
||||
response_class=JSONResponse
|
||||
)
|
||||
async def root():
|
||||
"""
|
||||
Get service information and health status
|
||||
|
||||
Returns basic information about the API service and available endpoints.
|
||||
"""
|
||||
logger.debug("Root endpoint accessed")
|
||||
return {
|
||||
"service": settings.app_name,
|
||||
"version": settings.app_version,
|
||||
"status": "healthy",
|
||||
"documentation": {
|
||||
"swagger_ui": "/docs",
|
||||
"redoc": "/redoc",
|
||||
"openapi_spec": "/openapi.json"
|
||||
},
|
||||
"endpoints": {
|
||||
"chat_completions": "/v1/chat/completions",
|
||||
"models": "/v1/models",
|
||||
"conversations": "/v1/conversations",
|
||||
"web_scraper": "/web-scraper/scrape",
|
||||
"health": "/health"
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# Health check endpoint
|
||||
@app.get(
|
||||
"/health",
|
||||
tags=["Health"],
|
||||
summary="Health check",
|
||||
response_class=JSONResponse
|
||||
)
|
||||
async def health_check():
|
||||
"""
|
||||
Simple health check endpoint for container orchestration
|
||||
|
||||
Returns a 200 OK status when the service is running properly.
|
||||
Used by Docker, Kubernetes, and load balancers.
|
||||
"""
|
||||
ollama_client = get_ollama_client()
|
||||
ollama_healthy = await ollama_client.health_check()
|
||||
|
||||
return {
|
||||
"status": "healthy",
|
||||
"ollama_connected": ollama_healthy
|
||||
}
|
||||
|
||||
|
||||
# Include routers
|
||||
app.include_router(chat_router) # /v1/chat/completions
|
||||
app.include_router(models_router) # /v1/models
|
||||
app.include_router(conversations_router) # /v1/conversations
|
||||
app.include_router(web_scraper_router) # /web-scraper/scrape
|
||||
|
||||
|
||||
# Global exception handler
|
||||
@app.exception_handler(Exception)
|
||||
async def global_exception_handler(request, exc):
|
||||
"""
|
||||
Catch-all exception handler for unhandled errors
|
||||
|
||||
Args:
|
||||
request: The request that caused the exception
|
||||
exc: The exception instance
|
||||
|
||||
Returns:
|
||||
JSON error response
|
||||
"""
|
||||
logger.error(f"Unhandled exception: {str(exc)}", exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={
|
||||
"detail": "Internal server error",
|
||||
"type": type(exc).__name__
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
Memory system for conversation persistence
|
||||
|
||||
Simplified architecture:
|
||||
- Tier 1: ConversationBufferMemory (in-memory, fast, last 10 turns)
|
||||
- Tier 2/3: QdrantConversationMemory (unified persistent + semantic search)
|
||||
- Manager: MemoryManager (orchestrates all tiers)
|
||||
"""
|
||||
from .tier1_buffer import ConversationBufferMemory, get_buffer_memory
|
||||
from .qdrant_memory import QdrantConversationMemory, get_qdrant_memory
|
||||
from .manager import MemoryManager, get_memory_manager
|
||||
from .schemas import (
|
||||
ConversationTurn,
|
||||
ConversationBuffer,
|
||||
ConversationMetadata,
|
||||
ConversationSummary,
|
||||
MemoryQuery,
|
||||
MemoryResult,
|
||||
MessageRole,
|
||||
TokenUsage
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Manager (primary interface)
|
||||
"MemoryManager",
|
||||
"get_memory_manager",
|
||||
# Tier 1
|
||||
"ConversationBufferMemory",
|
||||
"get_buffer_memory",
|
||||
# Tier 2/3
|
||||
"QdrantConversationMemory",
|
||||
"get_qdrant_memory",
|
||||
# Schemas
|
||||
"ConversationTurn",
|
||||
"ConversationBuffer",
|
||||
"ConversationMetadata",
|
||||
"ConversationSummary",
|
||||
"MemoryQuery",
|
||||
"MemoryResult",
|
||||
"MessageRole",
|
||||
"TokenUsage",
|
||||
]
|
||||
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
Base classes for memory system
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Optional
|
||||
from .schemas import ConversationTurn, ConversationBuffer, MemoryQuery, MemoryResult
|
||||
|
||||
|
||||
class BaseMemory(ABC):
|
||||
"""Base class for all memory tiers"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a new turn to memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns from memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists in this memory tier
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation exists
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier1Memory(BaseMemory):
|
||||
"""Base class for Tier 1 (working memory)"""
|
||||
|
||||
@abstractmethod
|
||||
async def get_buffer(self, conversation_id: str) -> Optional[ConversationBuffer]:
|
||||
"""
|
||||
Get the full conversation buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
ConversationBuffer or None if not found
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def prune(self, conversation_id: str, keep_last: int = 5) -> None:
|
||||
"""
|
||||
Prune old turns, keeping only the most recent ones
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
keep_last: Number of recent turns to keep
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier2Memory(BaseMemory):
|
||||
"""Base class for Tier 2 (short-term memory with summaries)"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_summary(
|
||||
self,
|
||||
conversation_id: str,
|
||||
summary_text: str,
|
||||
turn_range_start: int,
|
||||
turn_range_end: int
|
||||
) -> None:
|
||||
"""
|
||||
Add a conversation summary
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
summary_text: The summarized text
|
||||
turn_range_start: First turn number in summary
|
||||
turn_range_end: Last turn number in summary
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_summaries(self, conversation_id: str) -> List[dict]:
|
||||
"""
|
||||
Get all summaries for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
List of summary dictionaries
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Tier3Memory(BaseMemory):
|
||||
"""Base class for Tier 3 (long-term vector memory)"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_turn_with_embedding(
|
||||
self,
|
||||
conversation_id: str,
|
||||
turn: ConversationTurn,
|
||||
embedding: List[float]
|
||||
) -> None:
|
||||
"""
|
||||
Add a turn with its vector embedding
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn
|
||||
embedding: Vector embedding of the turn content
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def similarity_search(
|
||||
self,
|
||||
query_embedding: List[float],
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[dict]:
|
||||
"""
|
||||
Perform semantic similarity search
|
||||
|
||||
Args:
|
||||
query_embedding: Vector embedding of the search query
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,319 @@
|
||||
"""
|
||||
Memory Manager: Orchestrates all memory tiers
|
||||
|
||||
Coordinates:
|
||||
- Tier 1: ConversationBufferMemory (RAM, fast, last N turns)
|
||||
- Tier 2/3: QdrantConversationMemory (persistent + semantic)
|
||||
|
||||
Provides unified interface for memory operations with automatic
|
||||
tier management and consolidation.
|
||||
"""
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
from .tier1_buffer import ConversationBufferMemory, get_buffer_memory
|
||||
from .qdrant_memory import QdrantConversationMemory, get_qdrant_memory
|
||||
from .schemas import ConversationTurn, MessageRole, TokenUsage
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class MemoryManager:
|
||||
"""
|
||||
Unified memory manager orchestrating all tiers
|
||||
|
||||
Responsibilities:
|
||||
- Add turns to appropriate tiers
|
||||
- Retrieve conversation history (buffer + persistent)
|
||||
- Consolidate buffer to persistent storage
|
||||
- Semantic search across all conversations
|
||||
- Memory lifecycle management
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
buffer_memory: Optional[ConversationBufferMemory] = None,
|
||||
qdrant_memory: Optional[QdrantConversationMemory] = None,
|
||||
auto_consolidate: bool = True
|
||||
):
|
||||
"""
|
||||
Initialize memory manager
|
||||
|
||||
Args:
|
||||
buffer_memory: Optional Tier 1 buffer instance
|
||||
qdrant_memory: Optional Tier 2/3 Qdrant instance
|
||||
auto_consolidate: Automatically consolidate when buffer threshold reached
|
||||
"""
|
||||
self.buffer_memory = buffer_memory or get_buffer_memory()
|
||||
self.qdrant_memory = qdrant_memory or get_qdrant_memory()
|
||||
self.auto_consolidate = auto_consolidate
|
||||
|
||||
logger.info(
|
||||
f"MemoryManager initialized (auto_consolidate={auto_consolidate})"
|
||||
)
|
||||
|
||||
async def add_turn(
|
||||
self,
|
||||
conversation_id: str,
|
||||
role: MessageRole,
|
||||
content: str,
|
||||
tokens: Optional[TokenUsage] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
) -> ConversationTurn:
|
||||
"""
|
||||
Add a conversation turn to memory
|
||||
|
||||
Automatically:
|
||||
1. Adds to Tier 1 (buffer)
|
||||
2. Checks if consolidation threshold reached
|
||||
3. Consolidates to Tier 2/3 if needed
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
role: Message role (user, assistant, system)
|
||||
content: Message content
|
||||
tokens: Optional token usage
|
||||
metadata: Optional metadata
|
||||
|
||||
Returns:
|
||||
The created conversation turn
|
||||
"""
|
||||
# Get current buffer to determine turn number
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
turn_number = (buffer.metadata.turn_count + 1) if buffer else 1
|
||||
|
||||
# Create turn
|
||||
turn = ConversationTurn(
|
||||
role=role,
|
||||
content=content,
|
||||
timestamp=datetime.utcnow(),
|
||||
turn_number=turn_number,
|
||||
tokens=tokens,
|
||||
metadata=metadata or {}
|
||||
)
|
||||
|
||||
# Add to Tier 1 (buffer)
|
||||
await self.buffer_memory.add_turn(conversation_id, turn)
|
||||
logger.debug(f"Turn {turn_number} added to buffer for {conversation_id}")
|
||||
|
||||
# Check consolidation threshold
|
||||
if self.auto_consolidate:
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
if buffer.metadata.turn_count >= settings.memory_consolidation_threshold:
|
||||
logger.info(
|
||||
f"Consolidation threshold reached for {conversation_id} "
|
||||
f"({buffer.metadata.turn_count} turns)"
|
||||
)
|
||||
await self._consolidate_buffer(conversation_id)
|
||||
|
||||
return turn
|
||||
|
||||
async def get_recent_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 10
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get recent conversation turns (from buffer)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
|
||||
Returns:
|
||||
List of recent conversation turns
|
||||
"""
|
||||
return await self.buffer_memory.get_recent_turns(conversation_id, limit)
|
||||
|
||||
async def get_full_history(
|
||||
self,
|
||||
conversation_id: str,
|
||||
include_buffer: bool = True
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get complete conversation history
|
||||
|
||||
Combines:
|
||||
- Tier 2/3: Persistent history from Qdrant
|
||||
- Tier 1: Recent buffer (if include_buffer=True)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
include_buffer: Include buffer turns not yet consolidated
|
||||
|
||||
Returns:
|
||||
Complete conversation history, sorted chronologically
|
||||
"""
|
||||
# Get from Qdrant (Tier 2)
|
||||
qdrant_turns = await self.qdrant_memory.get_turns(conversation_id)
|
||||
|
||||
# Get from buffer (Tier 1)
|
||||
if include_buffer:
|
||||
buffer_turns = await self.buffer_memory.get_turns(conversation_id)
|
||||
|
||||
# Combine and deduplicate (Qdrant is source of truth)
|
||||
qdrant_turn_numbers = {t.turn_number for t in qdrant_turns}
|
||||
new_buffer_turns = [
|
||||
t for t in buffer_turns
|
||||
if t.turn_number not in qdrant_turn_numbers
|
||||
]
|
||||
|
||||
all_turns = qdrant_turns + new_buffer_turns
|
||||
else:
|
||||
all_turns = qdrant_turns
|
||||
|
||||
# Sort chronologically
|
||||
all_turns.sort(key=lambda t: t.turn_number)
|
||||
|
||||
return all_turns
|
||||
|
||||
async def search_conversations(
|
||||
self,
|
||||
query: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Semantic search across conversations (Tier 3 mode)
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
return await self.qdrant_memory.similarity_search(
|
||||
query=query,
|
||||
conversation_id=conversation_id,
|
||||
limit=limit
|
||||
)
|
||||
|
||||
async def consolidate(self, conversation_id: str) -> int:
|
||||
"""
|
||||
Manually trigger consolidation for a conversation
|
||||
|
||||
Moves all buffer turns to Qdrant (Tier 1 → Tier 2/3)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Number of turns consolidated
|
||||
"""
|
||||
return await self._consolidate_buffer(conversation_id)
|
||||
|
||||
async def _consolidate_buffer(self, conversation_id: str) -> int:
|
||||
"""
|
||||
Internal consolidation: Move buffer turns to Qdrant
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Number of turns consolidated
|
||||
"""
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
if not buffer or len(buffer.turns) == 0:
|
||||
logger.debug(f"No turns to consolidate for {conversation_id}")
|
||||
return 0
|
||||
|
||||
# Get turns from buffer
|
||||
buffer_turns = buffer.turns.copy()
|
||||
|
||||
# Add to Qdrant
|
||||
consolidated_count = 0
|
||||
for turn in buffer_turns:
|
||||
try:
|
||||
await self.qdrant_memory.add_turn(conversation_id, turn)
|
||||
consolidated_count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Error consolidating turn {turn.turn_number}: {e}")
|
||||
|
||||
logger.info(
|
||||
f"Consolidated {consolidated_count}/{len(buffer_turns)} turns "
|
||||
f"for {conversation_id}"
|
||||
)
|
||||
|
||||
# Note: We keep the buffer, just stored in Qdrant as well
|
||||
# Buffer will be pruned naturally as new turns come in
|
||||
# This provides redundancy and fast access to recent turns
|
||||
|
||||
return consolidated_count
|
||||
|
||||
async def clear_conversation(
|
||||
self,
|
||||
conversation_id: str,
|
||||
clear_buffer: bool = True,
|
||||
clear_qdrant: bool = True
|
||||
) -> None:
|
||||
"""
|
||||
Clear conversation from memory
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
clear_buffer: Clear from Tier 1 buffer
|
||||
clear_qdrant: Clear from Tier 2/3 Qdrant
|
||||
"""
|
||||
if clear_buffer:
|
||||
await self.buffer_memory.clear_conversation(conversation_id)
|
||||
logger.info(f"Cleared buffer for {conversation_id}")
|
||||
|
||||
if clear_qdrant:
|
||||
await self.qdrant_memory.clear_conversation(conversation_id)
|
||||
logger.info(f"Cleared Qdrant for {conversation_id}")
|
||||
|
||||
async def get_conversation_stats(
|
||||
self,
|
||||
conversation_id: str
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Get conversation statistics across all tiers
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Dictionary with stats from buffer and Qdrant
|
||||
"""
|
||||
# Get buffer stats
|
||||
buffer = await self.buffer_memory.get_buffer(conversation_id)
|
||||
buffer_stats = {
|
||||
"buffer_turns": buffer.metadata.turn_count if buffer else 0,
|
||||
"buffer_tokens": buffer.metadata.total_tokens if buffer else 0
|
||||
}
|
||||
|
||||
# Get Qdrant stats
|
||||
qdrant_stats = await self.qdrant_memory.get_conversation_stats(conversation_id)
|
||||
|
||||
# Combine
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
**buffer_stats,
|
||||
"qdrant_turns": qdrant_stats["total_turns"],
|
||||
"qdrant_tokens": qdrant_stats["total_tokens"],
|
||||
"exists_in_buffer": buffer is not None,
|
||||
"exists_in_qdrant": qdrant_stats["exists"]
|
||||
}
|
||||
|
||||
|
||||
# Global instance
|
||||
_memory_manager: Optional[MemoryManager] = None
|
||||
|
||||
|
||||
def get_memory_manager() -> MemoryManager:
|
||||
"""
|
||||
Get or create global memory manager instance
|
||||
|
||||
Returns:
|
||||
MemoryManager instance
|
||||
"""
|
||||
global _memory_manager
|
||||
if _memory_manager is None:
|
||||
_memory_manager = MemoryManager()
|
||||
return _memory_manager
|
||||
@@ -0,0 +1,387 @@
|
||||
"""
|
||||
Unified Tier 2/3: Qdrant-based conversation memory
|
||||
|
||||
Single Qdrant collection serving both purposes:
|
||||
- Tier 2: Historical retrieval (filter by conversation_id, time-based)
|
||||
- Tier 3: Semantic search (vector similarity across conversations)
|
||||
"""
|
||||
import logging
|
||||
import uuid
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import (
|
||||
Distance,
|
||||
VectorParams,
|
||||
PointStruct,
|
||||
Filter,
|
||||
FieldCondition,
|
||||
MatchValue,
|
||||
Range,
|
||||
)
|
||||
|
||||
from .base import BaseMemory
|
||||
from .schemas import ConversationTurn, MessageRole
|
||||
from src.config import get_settings
|
||||
from src.models.embeddings import get_embedding_client
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class QdrantConversationMemory(BaseMemory):
|
||||
"""
|
||||
Unified conversation memory using Qdrant
|
||||
|
||||
Stores all conversation turns with vectors for semantic search.
|
||||
Can be queried in two ways:
|
||||
- Tier 2 mode: Filter by conversation_id for chronological history
|
||||
- Tier 3 mode: Vector similarity search for semantic recall
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
host: Optional[str] = None,
|
||||
port: Optional[int] = None
|
||||
):
|
||||
"""
|
||||
Initialize Qdrant memory
|
||||
|
||||
Args:
|
||||
collection_name: Name of Qdrant collection
|
||||
host: Qdrant host
|
||||
port: Qdrant port
|
||||
"""
|
||||
self.collection_name = collection_name or settings.qdrant_collection_conversations
|
||||
self.host = host or settings.qdrant_host
|
||||
self.port = port or settings.qdrant_port
|
||||
|
||||
# Initialize clients
|
||||
self.client = QdrantClient(host=self.host, port=self.port)
|
||||
self.embedding_client = get_embedding_client()
|
||||
|
||||
logger.info(
|
||||
f"Initialized QdrantConversationMemory: "
|
||||
f"{self.host}:{self.port}/{self.collection_name}"
|
||||
)
|
||||
|
||||
# Ensure collection exists
|
||||
self._ensure_collection()
|
||||
|
||||
def _ensure_collection(self) -> None:
|
||||
"""Create collection if it doesn't exist"""
|
||||
try:
|
||||
collections = self.client.get_collections().collections
|
||||
collection_names = [c.name for c in collections]
|
||||
|
||||
if self.collection_name not in collection_names:
|
||||
logger.info(f"Creating collection: {self.collection_name}")
|
||||
self.client.create_collection(
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=VectorParams(
|
||||
size=settings.embedding_dimension,
|
||||
distance=Distance.COSINE
|
||||
)
|
||||
)
|
||||
logger.info(f"✓ Collection created: {self.collection_name}")
|
||||
else:
|
||||
logger.info(f"✓ Collection exists: {self.collection_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error ensuring collection: {e}")
|
||||
raise
|
||||
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a conversation turn with its embedding
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
# Generate embedding
|
||||
embedding = self.embedding_client.embed_text(turn.content)
|
||||
|
||||
# Create point ID: deterministic UUID from conversation_id + turn_number
|
||||
# Qdrant requires UUID or unsigned int, so we generate UUID from string
|
||||
point_id_str = f"{conversation_id}_{turn.turn_number}"
|
||||
point_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, point_id_str))
|
||||
|
||||
# Build payload
|
||||
payload = {
|
||||
"conversation_id": conversation_id,
|
||||
"turn_number": turn.turn_number,
|
||||
"role": turn.role.value if isinstance(turn.role, MessageRole) else turn.role,
|
||||
"content": turn.content,
|
||||
"timestamp": turn.timestamp.isoformat(),
|
||||
"metadata": turn.metadata,
|
||||
}
|
||||
|
||||
# Add token info if available
|
||||
if turn.tokens:
|
||||
payload["tokens_prompt"] = turn.tokens.prompt
|
||||
payload["tokens_completion"] = turn.tokens.completion
|
||||
payload["tokens_total"] = turn.tokens.total
|
||||
|
||||
# Upsert to Qdrant
|
||||
try:
|
||||
self.client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=[
|
||||
PointStruct(
|
||||
id=point_id,
|
||||
vector=embedding,
|
||||
payload=payload
|
||||
)
|
||||
]
|
||||
)
|
||||
logger.debug(f"Stored turn {turn.turn_number} for conversation {conversation_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing turn in Qdrant: {e}")
|
||||
raise
|
||||
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns for a conversation (Tier 2 mode: chronological)
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
try:
|
||||
# Scroll through all points for this conversation
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=limit or 100,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
with_vectors=False
|
||||
)
|
||||
|
||||
# Convert to ConversationTurn objects
|
||||
turns = []
|
||||
for point in points:
|
||||
payload = point.payload
|
||||
turn = ConversationTurn(
|
||||
role=MessageRole(payload["role"]),
|
||||
content=payload["content"],
|
||||
timestamp=datetime.fromisoformat(payload["timestamp"]),
|
||||
turn_number=payload["turn_number"],
|
||||
metadata=payload.get("metadata", {})
|
||||
)
|
||||
turns.append(turn)
|
||||
|
||||
# Sort by turn_number
|
||||
turns.sort(key=lambda t: t.turn_number)
|
||||
|
||||
return turns
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving turns from Qdrant: {e}")
|
||||
return []
|
||||
|
||||
async def similarity_search(
|
||||
self,
|
||||
query: str,
|
||||
conversation_id: Optional[str] = None,
|
||||
limit: int = 5
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Semantic search for relevant turns (Tier 3 mode: semantic)
|
||||
|
||||
Args:
|
||||
query: Search query text
|
||||
conversation_id: Optional filter to specific conversation
|
||||
limit: Maximum number of results
|
||||
|
||||
Returns:
|
||||
List of matching turns with scores
|
||||
"""
|
||||
try:
|
||||
# Generate query embedding
|
||||
query_embedding = self.embedding_client.embed_text(query)
|
||||
|
||||
# Build filter if conversation_id specified
|
||||
search_filter = None
|
||||
if conversation_id:
|
||||
search_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Search in Qdrant
|
||||
results = self.client.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_embedding,
|
||||
query_filter=search_filter,
|
||||
limit=limit,
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
# Convert results
|
||||
matches = []
|
||||
for result in results:
|
||||
payload = result.payload
|
||||
match = {
|
||||
"conversation_id": payload["conversation_id"],
|
||||
"turn_number": payload["turn_number"],
|
||||
"role": payload["role"],
|
||||
"content": payload["content"],
|
||||
"timestamp": payload["timestamp"],
|
||||
"score": result.score,
|
||||
}
|
||||
matches.append(match)
|
||||
|
||||
logger.debug(
|
||||
f"Semantic search found {len(matches)} matches for query: {query[:50]}..."
|
||||
)
|
||||
|
||||
return matches
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in semantic search: {e}")
|
||||
return []
|
||||
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
try:
|
||||
# Delete all points with this conversation_id
|
||||
self.client.delete(
|
||||
collection_name=self.collection_name,
|
||||
points_selector=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
logger.info(f"Cleared conversation {conversation_id} from Qdrant")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing conversation: {e}")
|
||||
raise
|
||||
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation has any turns
|
||||
"""
|
||||
try:
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=1,
|
||||
with_payload=False,
|
||||
with_vectors=False
|
||||
)
|
||||
return len(points) > 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking conversation existence: {e}")
|
||||
return False
|
||||
|
||||
async def get_conversation_stats(self, conversation_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get statistics about a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
Dictionary with stats
|
||||
"""
|
||||
try:
|
||||
points, _ = self.client.scroll(
|
||||
collection_name=self.collection_name,
|
||||
scroll_filter=Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="conversation_id",
|
||||
match=MatchValue(value=conversation_id)
|
||||
)
|
||||
]
|
||||
),
|
||||
limit=1000, # Get all points
|
||||
with_payload=True,
|
||||
with_vectors=False
|
||||
)
|
||||
|
||||
total_turns = len(points)
|
||||
total_tokens = sum(
|
||||
point.payload.get("tokens_total", 0) for point in points
|
||||
)
|
||||
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
"total_turns": total_turns,
|
||||
"total_tokens": total_tokens,
|
||||
"exists": total_turns > 0
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting conversation stats: {e}")
|
||||
return {
|
||||
"conversation_id": conversation_id,
|
||||
"total_turns": 0,
|
||||
"total_tokens": 0,
|
||||
"exists": False
|
||||
}
|
||||
|
||||
|
||||
# Global instance
|
||||
_qdrant_memory: Optional[QdrantConversationMemory] = None
|
||||
|
||||
|
||||
def get_qdrant_memory() -> QdrantConversationMemory:
|
||||
"""
|
||||
Get or create global Qdrant memory instance
|
||||
|
||||
Returns:
|
||||
QdrantConversationMemory instance
|
||||
"""
|
||||
global _qdrant_memory
|
||||
if _qdrant_memory is None:
|
||||
_qdrant_memory = QdrantConversationMemory()
|
||||
return _qdrant_memory
|
||||
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
Pydantic schemas for memory system
|
||||
"""
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MessageRole(str, Enum):
|
||||
"""Message role types"""
|
||||
SYSTEM = "system"
|
||||
USER = "user"
|
||||
ASSISTANT = "assistant"
|
||||
|
||||
|
||||
class TokenUsage(BaseModel):
|
||||
"""Token usage information"""
|
||||
prompt: int = 0
|
||||
completion: int = 0
|
||||
total: int = 0
|
||||
|
||||
|
||||
class ConversationTurn(BaseModel):
|
||||
"""A single turn in a conversation"""
|
||||
role: MessageRole
|
||||
content: str
|
||||
timestamp: datetime = Field(default_factory=datetime.utcnow)
|
||||
turn_number: int
|
||||
tokens: Optional[TokenUsage] = None
|
||||
metadata: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class ConversationMetadata(BaseModel):
|
||||
"""Metadata about a conversation"""
|
||||
conversation_id: str
|
||||
user_id: Optional[str] = None
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
last_updated: datetime = Field(default_factory=datetime.utcnow)
|
||||
turn_count: int = 0
|
||||
total_tokens: int = 0
|
||||
status: str = "active" # active, archived, deleted
|
||||
|
||||
|
||||
class ConversationBuffer(BaseModel):
|
||||
"""In-memory conversation buffer (Tier 1)"""
|
||||
conversation_id: str
|
||||
turns: List[ConversationTurn] = Field(default_factory=list)
|
||||
metadata: ConversationMetadata
|
||||
|
||||
|
||||
class ConversationSummary(BaseModel):
|
||||
"""Summarized conversation segment (Tier 2)"""
|
||||
conversation_id: str
|
||||
summary_text: str
|
||||
turn_range_start: int
|
||||
turn_range_end: int
|
||||
created_at: datetime = Field(default_factory=datetime.utcnow)
|
||||
token_count: int = 0
|
||||
|
||||
|
||||
class MemoryQuery(BaseModel):
|
||||
"""Query for memory retrieval"""
|
||||
conversation_id: str
|
||||
query: Optional[str] = None
|
||||
limit: int = Field(default=10, ge=1, le=100)
|
||||
include_tier1: bool = True
|
||||
include_tier2: bool = True
|
||||
include_tier3: bool = True
|
||||
|
||||
|
||||
class MemoryResult(BaseModel):
|
||||
"""Result from memory retrieval"""
|
||||
conversation_id: str
|
||||
turns: List[ConversationTurn] = Field(default_factory=list)
|
||||
summaries: List[ConversationSummary] = Field(default_factory=list)
|
||||
source_tiers: List[int] = Field(default_factory=list) # Which tiers contributed
|
||||
total_results: int = 0
|
||||
|
||||
|
||||
# API Request/Response Models
|
||||
|
||||
class ConversationListResponse(BaseModel):
|
||||
"""Response for listing conversations"""
|
||||
conversations: List[ConversationMetadata]
|
||||
total: int
|
||||
page: int = 1
|
||||
page_size: int = 50
|
||||
|
||||
|
||||
class ConversationDetailResponse(BaseModel):
|
||||
"""Response for conversation details"""
|
||||
metadata: ConversationMetadata
|
||||
recent_turns: List[ConversationTurn]
|
||||
turn_count: int
|
||||
|
||||
|
||||
class ConversationSearchRequest(BaseModel):
|
||||
"""Request for semantic search in conversation"""
|
||||
query: str
|
||||
limit: int = Field(default=5, ge=1, le=50)
|
||||
|
||||
|
||||
class ConversationSearchResponse(BaseModel):
|
||||
"""Response for semantic search"""
|
||||
conversation_id: str
|
||||
results: List[ConversationTurn]
|
||||
scores: List[float] = Field(default_factory=list)
|
||||
total_results: int
|
||||
@@ -0,0 +1,239 @@
|
||||
"""
|
||||
Tier 1: ConversationBufferMemory (In-Memory Working Memory)
|
||||
|
||||
Fast in-memory storage for recent conversation turns.
|
||||
- Stores last N turns in RAM
|
||||
- < 1ms access time
|
||||
- Ephemeral (lost on restart)
|
||||
- Automatic pruning when limit reached
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
from datetime import datetime
|
||||
from collections import OrderedDict
|
||||
|
||||
from .base import Tier1Memory
|
||||
from .schemas import (
|
||||
ConversationTurn,
|
||||
ConversationBuffer,
|
||||
ConversationMetadata,
|
||||
MessageRole,
|
||||
TokenUsage
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConversationBufferMemory(Tier1Memory):
|
||||
"""
|
||||
In-memory buffer for recent conversation turns.
|
||||
|
||||
Stores the last N turns of each conversation in RAM for fast access.
|
||||
Automatically prunes old turns when limit is reached.
|
||||
"""
|
||||
|
||||
def __init__(self, max_turns: int = 10):
|
||||
"""
|
||||
Initialize buffer memory
|
||||
|
||||
Args:
|
||||
max_turns: Maximum number of turns to keep per conversation
|
||||
"""
|
||||
self.max_turns = max_turns
|
||||
# Use OrderedDict to maintain insertion order
|
||||
self._buffers: Dict[str, ConversationBuffer] = OrderedDict()
|
||||
logger.info(f"Initialized ConversationBufferMemory with max_turns={max_turns}")
|
||||
|
||||
async def add_turn(self, conversation_id: str, turn: ConversationTurn) -> None:
|
||||
"""
|
||||
Add a new turn to the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
turn: The conversation turn to store
|
||||
"""
|
||||
# Get or create buffer
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
buffer = ConversationBuffer(
|
||||
conversation_id=conversation_id,
|
||||
turns=[],
|
||||
metadata=ConversationMetadata(
|
||||
conversation_id=conversation_id
|
||||
)
|
||||
)
|
||||
self._buffers[conversation_id] = buffer
|
||||
|
||||
# Add turn
|
||||
buffer.turns.append(turn)
|
||||
|
||||
# Update metadata
|
||||
buffer.metadata.turn_count = len(buffer.turns)
|
||||
buffer.metadata.last_updated = datetime.utcnow()
|
||||
|
||||
if turn.tokens:
|
||||
buffer.metadata.total_tokens += turn.tokens.total
|
||||
|
||||
# Auto-prune if exceeds max turns
|
||||
if len(buffer.turns) > self.max_turns:
|
||||
await self.prune(conversation_id, keep_last=self.max_turns)
|
||||
|
||||
logger.debug(
|
||||
f"Added turn {turn.turn_number} to conversation {conversation_id}. "
|
||||
f"Buffer size: {len(buffer.turns)}"
|
||||
)
|
||||
|
||||
async def get_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: Optional[int] = None,
|
||||
offset: int = 0
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Retrieve turns from the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Maximum number of turns to retrieve
|
||||
offset: Number of turns to skip
|
||||
|
||||
Returns:
|
||||
List of conversation turns
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return []
|
||||
|
||||
turns = buffer.turns[offset:]
|
||||
if limit:
|
||||
turns = turns[:limit]
|
||||
|
||||
return turns
|
||||
|
||||
async def get_recent_turns(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 10
|
||||
) -> List[ConversationTurn]:
|
||||
"""
|
||||
Get the most recent N turns
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
limit: Number of recent turns to retrieve
|
||||
|
||||
Returns:
|
||||
List of recent turns (most recent last)
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return []
|
||||
|
||||
return buffer.turns[-limit:] if len(buffer.turns) > limit else buffer.turns
|
||||
|
||||
async def get_buffer(self, conversation_id: str) -> Optional[ConversationBuffer]:
|
||||
"""
|
||||
Get the full conversation buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
ConversationBuffer or None if not found
|
||||
"""
|
||||
return self._buffers.get(conversation_id)
|
||||
|
||||
async def clear_conversation(self, conversation_id: str) -> None:
|
||||
"""
|
||||
Clear all turns for a conversation
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
"""
|
||||
if conversation_id in self._buffers:
|
||||
del self._buffers[conversation_id]
|
||||
logger.info(f"Cleared buffer for conversation {conversation_id}")
|
||||
|
||||
async def conversation_exists(self, conversation_id: str) -> bool:
|
||||
"""
|
||||
Check if a conversation exists in the buffer
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
|
||||
Returns:
|
||||
True if conversation exists
|
||||
"""
|
||||
return conversation_id in self._buffers
|
||||
|
||||
async def prune(self, conversation_id: str, keep_last: int = 5) -> None:
|
||||
"""
|
||||
Prune old turns, keeping only the most recent ones
|
||||
|
||||
Args:
|
||||
conversation_id: Unique conversation identifier
|
||||
keep_last: Number of recent turns to keep
|
||||
"""
|
||||
buffer = await self.get_buffer(conversation_id)
|
||||
if buffer is None:
|
||||
return
|
||||
|
||||
if len(buffer.turns) > keep_last:
|
||||
removed_count = len(buffer.turns) - keep_last
|
||||
buffer.turns = buffer.turns[-keep_last:]
|
||||
buffer.metadata.turn_count = len(buffer.turns)
|
||||
|
||||
logger.debug(
|
||||
f"Pruned {removed_count} turns from conversation {conversation_id}. "
|
||||
f"Kept last {keep_last} turns."
|
||||
)
|
||||
|
||||
async def get_all_conversation_ids(self) -> List[str]:
|
||||
"""
|
||||
Get list of all conversation IDs in memory
|
||||
|
||||
Returns:
|
||||
List of conversation IDs
|
||||
"""
|
||||
return list(self._buffers.keys())
|
||||
|
||||
async def get_buffer_stats(self) -> dict:
|
||||
"""
|
||||
Get statistics about buffer memory usage
|
||||
|
||||
Returns:
|
||||
Dictionary with stats
|
||||
"""
|
||||
total_conversations = len(self._buffers)
|
||||
total_turns = sum(len(buf.turns) for buf in self._buffers.values())
|
||||
total_tokens = sum(buf.metadata.total_tokens for buf in self._buffers.values())
|
||||
|
||||
return {
|
||||
"total_conversations": total_conversations,
|
||||
"total_turns": total_turns,
|
||||
"total_tokens": total_tokens,
|
||||
"max_turns_per_conversation": self.max_turns,
|
||||
"avg_turns_per_conversation": (
|
||||
total_turns / total_conversations if total_conversations > 0 else 0
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
# Global instance
|
||||
_buffer_memory: Optional[ConversationBufferMemory] = None
|
||||
|
||||
|
||||
def get_buffer_memory(max_turns: int = 10) -> ConversationBufferMemory:
|
||||
"""
|
||||
Get or create the global buffer memory instance
|
||||
|
||||
Args:
|
||||
max_turns: Maximum turns per conversation
|
||||
|
||||
Returns:
|
||||
ConversationBufferMemory instance
|
||||
"""
|
||||
global _buffer_memory
|
||||
if _buffer_memory is None:
|
||||
_buffer_memory = ConversationBufferMemory(max_turns=max_turns)
|
||||
return _buffer_memory
|
||||
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
Embedding model client for text vectorization
|
||||
|
||||
Uses sentence-transformers for generating embeddings.
|
||||
"""
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class EmbeddingClient:
|
||||
"""Client for generating text embeddings"""
|
||||
|
||||
def __init__(self, model_name: Optional[str] = None):
|
||||
"""
|
||||
Initialize embedding client
|
||||
|
||||
Args:
|
||||
model_name: Optional model name, defaults to config
|
||||
"""
|
||||
self.model_name = model_name or settings.embedding_model
|
||||
self.dimension = settings.embedding_dimension
|
||||
self._model: Optional[SentenceTransformer] = None
|
||||
logger.info(f"Initializing EmbeddingClient with model: {self.model_name}")
|
||||
|
||||
def _load_model(self) -> SentenceTransformer:
|
||||
"""
|
||||
Lazy load the embedding model
|
||||
|
||||
Returns:
|
||||
Loaded SentenceTransformer model
|
||||
"""
|
||||
if self._model is None:
|
||||
logger.info(f"Loading embedding model: {self.model_name}")
|
||||
self._model = SentenceTransformer(self.model_name)
|
||||
logger.info(f"Model loaded successfully. Embedding dimension: {self.dimension}")
|
||||
return self._model
|
||||
|
||||
def embed_text(self, text: str) -> List[float]:
|
||||
"""
|
||||
Generate embedding for a single text
|
||||
|
||||
Args:
|
||||
text: Input text to embed
|
||||
|
||||
Returns:
|
||||
List of floats representing the embedding vector
|
||||
"""
|
||||
model = self._load_model()
|
||||
embedding = model.encode(text, convert_to_numpy=True)
|
||||
return embedding.tolist()
|
||||
|
||||
def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts
|
||||
|
||||
Args:
|
||||
texts: List of input texts
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
model = self._load_model()
|
||||
embeddings = model.encode(
|
||||
texts,
|
||||
batch_size=settings.embedding_batch_size,
|
||||
convert_to_numpy=True,
|
||||
show_progress_bar=False
|
||||
)
|
||||
return embeddings.tolist()
|
||||
|
||||
def get_dimension(self) -> int:
|
||||
"""
|
||||
Get embedding dimension
|
||||
|
||||
Returns:
|
||||
Embedding vector dimension
|
||||
"""
|
||||
return self.dimension
|
||||
|
||||
|
||||
# Global instance
|
||||
_embedding_client: Optional[EmbeddingClient] = None
|
||||
|
||||
|
||||
def get_embedding_client() -> EmbeddingClient:
|
||||
"""
|
||||
Get or create global embedding client instance
|
||||
|
||||
Returns:
|
||||
EmbeddingClient instance
|
||||
"""
|
||||
global _embedding_client
|
||||
if _embedding_client is None:
|
||||
_embedding_client = EmbeddingClient()
|
||||
return _embedding_client
|
||||
|
||||
|
||||
async def embed_text_async(text: str) -> List[float]:
|
||||
"""
|
||||
Async wrapper for embedding text
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
|
||||
Returns:
|
||||
Embedding vector
|
||||
"""
|
||||
client = get_embedding_client()
|
||||
return client.embed_text(text)
|
||||
|
||||
|
||||
async def embed_batch_async(texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Async wrapper for batch embedding
|
||||
|
||||
Args:
|
||||
texts: List of input texts
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
client = get_embedding_client()
|
||||
return client.embed_batch(texts)
|
||||
@@ -0,0 +1,201 @@
|
||||
"""
|
||||
Ollama client for model inference.
|
||||
Handles both streaming and non-streaming requests.
|
||||
"""
|
||||
|
||||
import httpx
|
||||
import json
|
||||
import logging
|
||||
from typing import AsyncIterator, Dict, Any, Optional
|
||||
from src.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class OllamaClient:
|
||||
"""Client for interacting with Ollama API."""
|
||||
|
||||
def __init__(self):
|
||||
self.base_url = settings.ollama_base_url
|
||||
self.timeout = settings.ollama_timeout
|
||||
self.client = httpx.AsyncClient(timeout=self.timeout)
|
||||
logger.info(f"Initialized Ollama client: {self.base_url}")
|
||||
|
||||
async def close(self):
|
||||
"""Close the HTTP client."""
|
||||
await self.client.aclose()
|
||||
|
||||
def resolve_model(self, model_name: str) -> str:
|
||||
"""
|
||||
Resolve model alias to actual Ollama model.
|
||||
|
||||
Args:
|
||||
model_name: Requested model name (e.g., "gpt-3.5-turbo")
|
||||
|
||||
Returns:
|
||||
Actual Ollama model name (e.g., "gemma:7b")
|
||||
"""
|
||||
resolved = settings.model_aliases.get(model_name, model_name)
|
||||
if resolved != model_name:
|
||||
logger.info(f"Model resolution: {model_name} → {resolved}")
|
||||
return resolved
|
||||
|
||||
async def generate_non_streaming(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Generate non-streaming response from Ollama.
|
||||
|
||||
Args:
|
||||
model: Model name
|
||||
prompt: User prompt
|
||||
temperature: Sampling temperature
|
||||
max_tokens: Maximum tokens to generate
|
||||
|
||||
Returns:
|
||||
Dict with 'response' and 'tokens' keys
|
||||
"""
|
||||
actual_model = self.resolve_model(model)
|
||||
|
||||
payload = {
|
||||
"model": actual_model,
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
"options": {
|
||||
"temperature": temperature,
|
||||
}
|
||||
}
|
||||
|
||||
if max_tokens:
|
||||
payload["options"]["num_predict"] = max_tokens
|
||||
|
||||
logger.debug(f"Ollama request to {actual_model}")
|
||||
|
||||
try:
|
||||
response = await self.client.post(
|
||||
f"{self.base_url}/api/generate",
|
||||
json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
return {
|
||||
"response": result.get("response", ""),
|
||||
"tokens": {
|
||||
"prompt": result.get("prompt_eval_count", 0),
|
||||
"completion": result.get("eval_count", 0),
|
||||
"total": result.get("prompt_eval_count", 0) + result.get("eval_count", 0)
|
||||
}
|
||||
}
|
||||
|
||||
except httpx.HTTPError as e:
|
||||
logger.error(f"Ollama request failed: {e}")
|
||||
raise
|
||||
|
||||
async def generate_streaming(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Generate streaming response from Ollama.
|
||||
|
||||
Args:
|
||||
model: Model name
|
||||
prompt: User prompt
|
||||
temperature: Sampling temperature
|
||||
max_tokens: Maximum tokens to generate
|
||||
|
||||
Yields:
|
||||
Token strings
|
||||
"""
|
||||
actual_model = self.resolve_model(model)
|
||||
|
||||
payload = {
|
||||
"model": actual_model,
|
||||
"prompt": prompt,
|
||||
"stream": True,
|
||||
"options": {
|
||||
"temperature": temperature,
|
||||
}
|
||||
}
|
||||
|
||||
if max_tokens:
|
||||
payload["options"]["num_predict"] = max_tokens
|
||||
|
||||
logger.debug(f"Ollama streaming request to {actual_model}")
|
||||
|
||||
try:
|
||||
async with self.client.stream(
|
||||
"POST",
|
||||
f"{self.base_url}/api/generate",
|
||||
json=payload
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
chunk = json.loads(line)
|
||||
if "response" in chunk:
|
||||
token = chunk["response"]
|
||||
if token:
|
||||
yield token
|
||||
|
||||
# Check if done
|
||||
if chunk.get("done", False):
|
||||
break
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse JSON: {line}")
|
||||
continue
|
||||
|
||||
except httpx.HTTPError as e:
|
||||
logger.error(f"Ollama streaming request failed: {e}")
|
||||
raise
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""
|
||||
Check if Ollama is healthy.
|
||||
|
||||
Returns:
|
||||
True if healthy, False otherwise
|
||||
"""
|
||||
try:
|
||||
response = await self.client.get(
|
||||
f"{self.base_url}/api/tags",
|
||||
timeout=5.0
|
||||
)
|
||||
return response.status_code == 200
|
||||
except Exception as e:
|
||||
logger.error(f"Ollama health check failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
# Global client instance
|
||||
_ollama_client: Optional[OllamaClient] = None
|
||||
|
||||
|
||||
def get_ollama_client() -> OllamaClient:
|
||||
"""Get or create the global Ollama client instance."""
|
||||
global _ollama_client
|
||||
if _ollama_client is None:
|
||||
_ollama_client = OllamaClient()
|
||||
return _ollama_client
|
||||
|
||||
|
||||
async def close_ollama_client():
|
||||
"""Close the global Ollama client."""
|
||||
global _ollama_client
|
||||
if _ollama_client is not None:
|
||||
await _ollama_client.close()
|
||||
_ollama_client = None
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
Web scraper module for extracting content from websites
|
||||
"""
|
||||
from src.web_scraper.router import router
|
||||
from src.web_scraper.schemas import WebScraperRequest, WebScraperResponse
|
||||
from src.web_scraper.service import WebScraperService
|
||||
|
||||
__all__ = [
|
||||
"router",
|
||||
"WebScraperRequest",
|
||||
"WebScraperResponse",
|
||||
"WebScraperService",
|
||||
]
|
||||
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
Configuration for web scraper module
|
||||
"""
|
||||
from pydantic_settings import BaseSettings
|
||||
from functools import lru_cache
|
||||
|
||||
|
||||
class WebScraperSettings(BaseSettings):
|
||||
"""Web scraper specific settings"""
|
||||
|
||||
# HTTP client configuration
|
||||
request_timeout: int = 30
|
||||
max_redirects: int = 5
|
||||
user_agent: str = "Mozilla/5.0 (compatible; CoreCode/1.0)"
|
||||
|
||||
# Content extraction
|
||||
default_max_length: int = 10000
|
||||
max_links_to_extract: int = 50
|
||||
|
||||
# Rate limiting (future use)
|
||||
rate_limit_enabled: bool = False
|
||||
requests_per_minute: int = 60
|
||||
|
||||
class Config:
|
||||
env_prefix = "WEB_SCRAPER_"
|
||||
case_sensitive = False
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def get_web_scraper_settings() -> WebScraperSettings:
|
||||
"""Cached web scraper settings instance"""
|
||||
return WebScraperSettings()
|
||||
@@ -0,0 +1,18 @@
|
||||
"""
|
||||
Custom exceptions for web scraper module
|
||||
"""
|
||||
|
||||
|
||||
class WebScraperException(Exception):
|
||||
"""Base exception for web scraper module"""
|
||||
pass
|
||||
|
||||
|
||||
class FetchError(WebScraperException):
|
||||
"""Raised when URL fetch fails"""
|
||||
pass
|
||||
|
||||
|
||||
class ScrapingError(WebScraperException):
|
||||
"""Raised when content extraction fails"""
|
||||
pass
|
||||
@@ -0,0 +1,79 @@
|
||||
"""
|
||||
API routes for web scraper module
|
||||
"""
|
||||
from fastapi import APIRouter, HTTPException, status
|
||||
|
||||
from src.logging_config import get_logger
|
||||
from src.web_scraper.schemas import WebScraperRequest, WebScraperResponse
|
||||
from src.web_scraper.service import WebScraperService
|
||||
from src.web_scraper.exceptions import FetchError, ScrapingError
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/web-scraper",
|
||||
tags=["Web Scraper"]
|
||||
)
|
||||
|
||||
# Initialize service (could be dependency injected for testing)
|
||||
scraper_service = WebScraperService()
|
||||
|
||||
|
||||
@router.post(
|
||||
"/scrape",
|
||||
response_model=WebScraperResponse,
|
||||
status_code=status.HTTP_200_OK,
|
||||
summary="Scrape website content",
|
||||
description="""
|
||||
Scrape and extract main content from a website.
|
||||
|
||||
Uses trafilatura for intelligent content extraction (articles, blog posts, documentation),
|
||||
with BeautifulSoup as fallback. Perfect for feeding webpage content to LLMs.
|
||||
|
||||
**Features:**
|
||||
- Intelligent main content extraction
|
||||
- Removes navigation, ads, footers
|
||||
- Optional link extraction
|
||||
- Configurable content length limits
|
||||
|
||||
**Rate Limiting:** None (internal network use only)
|
||||
"""
|
||||
)
|
||||
async def scrape_website(request: WebScraperRequest) -> WebScraperResponse:
|
||||
"""
|
||||
Scrape a website and extract its main content
|
||||
|
||||
Args:
|
||||
request: Scraping request with URL and options
|
||||
|
||||
Returns:
|
||||
Extracted content with metadata
|
||||
|
||||
Raises:
|
||||
HTTPException: 400 for fetch errors, 500 for processing errors
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Received scrape request for: {request.url}")
|
||||
result = await scraper_service.scrape_url(request)
|
||||
return result
|
||||
|
||||
except FetchError as e:
|
||||
logger.warning(f"Fetch failed: {str(e)}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Failed to fetch URL: {str(e)}"
|
||||
)
|
||||
|
||||
except ScrapingError as e:
|
||||
logger.error(f"Scraping failed: {str(e)}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Failed to extract content: {str(e)}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error: {str(e)}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="An unexpected error occurred"
|
||||
)
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
Pydantic schemas for web scraper module
|
||||
"""
|
||||
from pydantic import HttpUrl, Field
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
from src.base_schema import BaseSchema
|
||||
|
||||
|
||||
class WebScraperRequest(BaseSchema):
|
||||
"""Request model for web scraping"""
|
||||
|
||||
url: HttpUrl = Field(
|
||||
...,
|
||||
description="The URL to scrape",
|
||||
examples=["https://example.com/article"]
|
||||
)
|
||||
|
||||
extract_main_content: bool = Field(
|
||||
default=True,
|
||||
description="Use intelligent content extraction (trafilatura) vs raw HTML parsing"
|
||||
)
|
||||
|
||||
include_links: bool = Field(
|
||||
default=False,
|
||||
description="Include list of links found on the page"
|
||||
)
|
||||
|
||||
max_length: Optional[int] = Field(
|
||||
default=10000,
|
||||
ge=100,
|
||||
le=100000,
|
||||
description="Maximum content length to return (100-100000 chars)"
|
||||
)
|
||||
|
||||
|
||||
class WebScraperResponse(BaseSchema):
|
||||
"""Response model for web scraping"""
|
||||
|
||||
url: str = Field(
|
||||
...,
|
||||
description="The scraped URL"
|
||||
)
|
||||
|
||||
title: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Page title extracted from <title> tag"
|
||||
)
|
||||
|
||||
content: str = Field(
|
||||
...,
|
||||
description="Extracted page content"
|
||||
)
|
||||
|
||||
extracted_at: datetime = Field(
|
||||
...,
|
||||
description="UTC timestamp when content was extracted"
|
||||
)
|
||||
|
||||
content_length: int = Field(
|
||||
...,
|
||||
ge=0,
|
||||
description="Length of extracted content in characters"
|
||||
)
|
||||
|
||||
links: Optional[list[str]] = Field(
|
||||
default=None,
|
||||
description="List of HTTP(S) links found on the page (max 50)"
|
||||
)
|
||||
@@ -0,0 +1,213 @@
|
||||
"""
|
||||
Business logic for web scraper module
|
||||
"""
|
||||
import httpx
|
||||
from bs4 import BeautifulSoup
|
||||
import trafilatura
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
from src.logging_config import get_logger
|
||||
from src.web_scraper.config import get_web_scraper_settings
|
||||
from src.web_scraper.schemas import WebScraperRequest, WebScraperResponse
|
||||
from src.web_scraper.exceptions import ScrapingError, FetchError
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class WebScraperService:
|
||||
"""Service class for web scraping operations"""
|
||||
|
||||
def __init__(self):
|
||||
self.settings = get_web_scraper_settings()
|
||||
|
||||
async def scrape_url(self, request: WebScraperRequest) -> WebScraperResponse:
|
||||
"""
|
||||
Scrape and extract content from a URL
|
||||
|
||||
Args:
|
||||
request: Scraping request parameters
|
||||
|
||||
Returns:
|
||||
Extracted content with metadata
|
||||
|
||||
Raises:
|
||||
FetchError: If URL cannot be fetched
|
||||
ScrapingError: If content extraction fails
|
||||
"""
|
||||
url_str = str(request.url)
|
||||
logger.info(f"Starting scrape for URL: {url_str}")
|
||||
|
||||
try:
|
||||
# Fetch the webpage
|
||||
html_content = await self._fetch_url(url_str)
|
||||
|
||||
# Extract content based on settings
|
||||
if request.extract_main_content:
|
||||
content = self._extract_main_content(html_content, request.include_links)
|
||||
else:
|
||||
content = self._extract_basic_content(html_content)
|
||||
|
||||
# Extract metadata
|
||||
title = self._extract_title(html_content)
|
||||
links = self._extract_links(html_content) if request.include_links else None
|
||||
|
||||
# Clean and truncate content
|
||||
content = self._clean_content(content)
|
||||
if request.max_length and len(content) > request.max_length:
|
||||
content = content[:request.max_length] + "\n\n[Content truncated...]"
|
||||
logger.debug(f"Content truncated to {request.max_length} characters")
|
||||
|
||||
logger.info(f"Successfully scraped {len(content)} characters from {url_str}")
|
||||
|
||||
return WebScraperResponse(
|
||||
url=url_str,
|
||||
title=title,
|
||||
content=content,
|
||||
extracted_at=datetime.now(timezone.utc),
|
||||
content_length=len(content),
|
||||
links=links
|
||||
)
|
||||
|
||||
except FetchError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Scraping failed for {url_str}: {str(e)}", exc_info=True)
|
||||
raise ScrapingError(f"Failed to scrape content: {str(e)}")
|
||||
|
||||
async def _fetch_url(self, url: str) -> str:
|
||||
"""
|
||||
Fetch HTML content from URL
|
||||
|
||||
Args:
|
||||
url: URL to fetch
|
||||
|
||||
Returns:
|
||||
HTML content as string
|
||||
|
||||
Raises:
|
||||
FetchError: If fetch fails
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=self.settings.request_timeout,
|
||||
follow_redirects=True,
|
||||
max_redirects=self.settings.max_redirects
|
||||
) as client:
|
||||
logger.debug(f"Fetching URL: {url}")
|
||||
response = await client.get(
|
||||
url,
|
||||
headers={"User-Agent": self.settings.user_agent}
|
||||
)
|
||||
response.raise_for_status()
|
||||
logger.debug(f"Fetched {len(response.text)} bytes from {url}")
|
||||
return response.text
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.error(f"HTTP error {e.response.status_code} for {url}")
|
||||
raise FetchError(f"HTTP {e.response.status_code}: {e.response.reason_phrase}")
|
||||
except httpx.RequestError as e:
|
||||
logger.error(f"Request error for {url}: {str(e)}")
|
||||
raise FetchError(f"Failed to fetch URL: {str(e)}")
|
||||
|
||||
def _extract_main_content(self, html: str, include_links: bool = False) -> str:
|
||||
"""
|
||||
Extract main content using trafilatura (intelligent extraction)
|
||||
|
||||
Args:
|
||||
html: Raw HTML content
|
||||
include_links: Whether to preserve links in output
|
||||
|
||||
Returns:
|
||||
Extracted content
|
||||
"""
|
||||
logger.debug("Extracting main content with trafilatura")
|
||||
content = trafilatura.extract(
|
||||
html,
|
||||
include_links=include_links,
|
||||
include_images=False,
|
||||
output_format='txt',
|
||||
no_fallback=False
|
||||
)
|
||||
|
||||
# Fallback to BeautifulSoup if trafilatura fails
|
||||
if not content:
|
||||
logger.debug("Trafilatura extraction failed, falling back to BeautifulSoup")
|
||||
content = self._extract_basic_content(html)
|
||||
|
||||
return content
|
||||
|
||||
def _extract_basic_content(self, html: str) -> str:
|
||||
"""
|
||||
Extract content using basic BeautifulSoup parsing
|
||||
|
||||
Args:
|
||||
html: Raw HTML content
|
||||
|
||||
Returns:
|
||||
Extracted text content
|
||||
"""
|
||||
logger.debug("Extracting content with BeautifulSoup")
|
||||
soup = BeautifulSoup(html, 'html.parser')
|
||||
|
||||
# Remove unwanted elements
|
||||
for element in soup(["script", "style", "nav", "footer", "header", "aside"]):
|
||||
element.decompose()
|
||||
|
||||
# Extract text
|
||||
text = soup.get_text(separator='\n', strip=True)
|
||||
return text
|
||||
|
||||
def _extract_title(self, html: str) -> Optional[str]:
|
||||
"""
|
||||
Extract page title from HTML
|
||||
|
||||
Args:
|
||||
html: Raw HTML content
|
||||
|
||||
Returns:
|
||||
Page title or None
|
||||
"""
|
||||
soup = BeautifulSoup(html, 'html.parser')
|
||||
title = soup.title.string if soup.title else None
|
||||
if title:
|
||||
title = title.strip()
|
||||
logger.debug(f"Extracted title: {title}")
|
||||
return title
|
||||
|
||||
def _extract_links(self, html: str) -> list[str]:
|
||||
"""
|
||||
Extract HTTP(S) links from HTML
|
||||
|
||||
Args:
|
||||
html: Raw HTML content
|
||||
|
||||
Returns:
|
||||
List of absolute HTTP(S) URLs
|
||||
"""
|
||||
soup = BeautifulSoup(html, 'html.parser')
|
||||
links = [
|
||||
a.get('href')
|
||||
for a in soup.find_all('a', href=True)
|
||||
if a.get('href', '').startswith('http')
|
||||
]
|
||||
|
||||
# Limit number of links
|
||||
links = links[:self.settings.max_links_to_extract]
|
||||
logger.debug(f"Extracted {len(links)} links")
|
||||
return links
|
||||
|
||||
def _clean_content(self, content: str) -> str:
|
||||
"""
|
||||
Clean and normalize extracted content
|
||||
|
||||
Args:
|
||||
content: Raw extracted content
|
||||
|
||||
Returns:
|
||||
Cleaned content
|
||||
"""
|
||||
# Remove empty lines and normalize whitespace
|
||||
lines = [line.strip() for line in content.split('\n') if line.strip()]
|
||||
cleaned = '\n'.join(lines)
|
||||
return cleaned
|
||||
Reference in New Issue
Block a user