Spaces:
Running
Running
File size: 4,117 Bytes
5b89d45 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """
Pydantic schemas for FastAPI request/response models
"""
from pydantic import BaseModel, Field
from typing import Optional, List
from enum import Enum
class ProviderEnum(str, Enum):
gemini = "gemini"
groq = "groq"
class VectorDBEnum(str, Enum):
chroma = "chroma"
faiss = "faiss"
qdrant = "qdrant"
# ============================================================================
# Chat Schemas
# ============================================================================
class ChatRequest(BaseModel):
"""Request body for chat endpoint"""
question: str = Field(..., description="The question to ask about the codebase")
use_agent: bool = Field(default=True, description="Use agentic mode with tool calls")
provider: ProviderEnum = Field(default=ProviderEnum.gemini, description="LLM provider")
class Config:
json_schema_extra = {
"example": {
"question": "What does this codebase do?",
"use_agent": True,
"provider": "gemini"
}
}
class SourceInfo(BaseModel):
"""Information about a source file used in the response"""
file_path: str
relevance_score: Optional[float] = None
class ChatResponse(BaseModel):
"""Response from chat endpoint"""
answer: str = Field(..., description="The generated answer")
sources: List[SourceInfo] = Field(default=[], description="Source files used")
mode: str = Field(..., description="Mode used: 'agent' or 'linear'")
processing_time: float = Field(..., description="Time taken in seconds")
class Config:
json_schema_extra = {
"example": {
"answer": "This codebase implements a RAG-based code chatbot...",
"sources": [{"file_path": "code_chatbot/rag.py", "relevance_score": 0.95}],
"mode": "agent",
"processing_time": 2.5
}
}
# ============================================================================
# Index Schemas
# ============================================================================
class IndexRequest(BaseModel):
"""Request body for index endpoint"""
source: str = Field(..., description="GitHub URL, local path, or ZIP file path")
provider: ProviderEnum = Field(default=ProviderEnum.gemini, description="Embedding provider")
vector_db: VectorDBEnum = Field(default=VectorDBEnum.chroma, description="Vector database type")
class Config:
json_schema_extra = {
"example": {
"source": "https://github.com/user/repo",
"provider": "gemini",
"vector_db": "chroma"
}
}
class IndexResponse(BaseModel):
"""Response from index endpoint"""
status: str = Field(..., description="'success' or 'error'")
message: str = Field(..., description="Status message")
files_indexed: int = Field(default=0, description="Number of files indexed")
chunks_created: int = Field(default=0, description="Number of chunks created")
graph_nodes: int = Field(default=0, description="Number of AST graph nodes")
class Config:
json_schema_extra = {
"example": {
"status": "success",
"message": "Successfully indexed repository",
"files_indexed": 45,
"chunks_created": 1200,
"graph_nodes": 350
}
}
# ============================================================================
# Health Schemas
# ============================================================================
class HealthResponse(BaseModel):
"""Health check response"""
status: str = Field(..., description="'healthy' or 'unhealthy'")
indexed: bool = Field(..., description="Whether a codebase is currently indexed")
provider: Optional[str] = Field(None, description="Current LLM provider")
vector_db: Optional[str] = Field(None, description="Current vector database")
documents_count: int = Field(default=0, description="Number of indexed documents")
|