95 lines
2.3 KiB
Python
95 lines
2.3 KiB
Python
"""Pydantic v2 models for the RAG domain."""
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from enum import StrEnum
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
|
|
class EmbeddingProvider(StrEnum):
|
|
"""Embedding provider options."""
|
|
|
|
OLLAMA_BGE_M3 = "ollama_bge_m3"
|
|
OPENAI = "openai"
|
|
OPENROUTER = "openrouter"
|
|
HUGGINGFACE = "huggingface"
|
|
COHERE = "cohere"
|
|
|
|
|
|
# Default RAG collections (matches legacy app.rag_engine.COLLECTIONS)
|
|
COLLECTIONS: list[str] = [
|
|
"scam_intel",
|
|
"deployer_history",
|
|
"wallet_labels",
|
|
"contract_audit",
|
|
"phishing_db",
|
|
]
|
|
|
|
|
|
class SearchRequest(BaseModel):
|
|
"""RAG search request."""
|
|
|
|
model_config = ConfigDict(str_strip_whitespace=True)
|
|
|
|
query: str = Field(..., min_length=1, max_length=2048)
|
|
collection: str = Field(default="scam_intel")
|
|
top_k: int = Field(default=5, ge=1, le=50)
|
|
min_similarity: float = Field(default=0.0, ge=0.0, le=1.0)
|
|
filters: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class SearchHit(BaseModel):
|
|
"""A single search result."""
|
|
|
|
content: str
|
|
score: float = 0.0
|
|
metadata: dict[str, Any] = Field(default_factory=dict)
|
|
collection: str = ""
|
|
doc_id: str = ""
|
|
|
|
|
|
class SearchResponse(BaseModel):
|
|
"""RAG search response."""
|
|
|
|
query: str
|
|
hits: list[SearchHit] = Field(default_factory=list)
|
|
total: int = 0
|
|
took_ms: int = 0
|
|
collection: str = ""
|
|
|
|
|
|
class IngestRequest(BaseModel):
|
|
"""RAG document ingestion request."""
|
|
|
|
model_config = ConfigDict(str_strip_whitespace=True)
|
|
|
|
collection: str = Field(default="scam_intel")
|
|
content: str = Field(..., min_length=1)
|
|
doc_id: str | None = None
|
|
metadata: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class IngestResult(BaseModel):
|
|
"""RAG document ingestion result."""
|
|
|
|
doc_id: str
|
|
collection: str
|
|
status: str = "ok" # ok | failed
|
|
chunks: int = 0
|
|
error: str | None = None
|
|
ingested_at: datetime = Field(default_factory=datetime.utcnow)
|
|
|
|
|
|
class FeedbackRecord(BaseModel):
|
|
"""Scanner → RAG feedback record."""
|
|
|
|
model_config = ConfigDict(str_strip_whitespace=True)
|
|
|
|
token_address: str
|
|
chain: str = "solana"
|
|
safety_score: float
|
|
risk_flags: list[str] = Field(default_factory=list)
|
|
action: str = "ingest" # ingest | remove | update
|
|
source: str = "scanner"
|