File size: 2,377 Bytes
6993919 | 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 | """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"
|