Initial commit: breakpilot-core - Shared Infrastructure
Docker Compose with 24+ services: - PostgreSQL (PostGIS), Valkey, MinIO, Qdrant - Vault (PKI/TLS), Nginx (Reverse Proxy) - Backend Core API, Consent Service, Billing Service - RAG Service, Embedding Service - Gitea, Woodpecker CI/CD - Night Scheduler, Health Aggregator - Jitsi (Web/XMPP/JVB/Jicofo), Mailpit Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
123
rag-service/embedding_client.py
Normal file
123
rag-service/embedding_client.py
Normal file
@@ -0,0 +1,123 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from config import settings
|
||||
|
||||
logger = logging.getLogger("rag-service.embedding")
|
||||
|
||||
_TIMEOUT = httpx.Timeout(timeout=120.0, connect=10.0)
|
||||
|
||||
|
||||
class EmbeddingClient:
|
||||
"""HTTP client for the embedding-service (port 8087)."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._base_url: str = settings.EMBEDDING_SERVICE_URL.rstrip("/")
|
||||
|
||||
def _url(self, path: str) -> str:
|
||||
return f"{self._base_url}{path}"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Embeddings
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def generate_embeddings(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Send a batch of texts to the embedding service and return a list of
|
||||
embedding vectors.
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
self._url("/api/v1/embeddings"),
|
||||
json={"texts": texts},
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data.get("embeddings", [])
|
||||
|
||||
async def generate_single_embedding(self, text: str) -> list[float]:
|
||||
"""Convenience wrapper for a single text."""
|
||||
results = await self.generate_embeddings([text])
|
||||
if not results:
|
||||
raise ValueError("Embedding service returned empty result")
|
||||
return results[0]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Reranking
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def rerank_documents(
|
||||
self,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
top_k: int = 10,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Ask the embedding service to re-rank documents for a given query.
|
||||
Returns a list of {index, score, text}.
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
self._url("/api/v1/rerank"),
|
||||
json={
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
"top_k": top_k,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data.get("results", [])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Chunking
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def chunk_text(
|
||||
self,
|
||||
text: str,
|
||||
strategy: str = "recursive",
|
||||
chunk_size: int = 512,
|
||||
overlap: int = 50,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Ask the embedding service to chunk a long text.
|
||||
Returns a list of chunk strings.
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
self._url("/api/v1/chunk"),
|
||||
json={
|
||||
"text": text,
|
||||
"strategy": strategy,
|
||||
"chunk_size": chunk_size,
|
||||
"overlap": overlap,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data.get("chunks", [])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# PDF extraction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def extract_pdf(self, pdf_bytes: bytes) -> str:
|
||||
"""
|
||||
Send raw PDF bytes to the embedding service for text extraction.
|
||||
Returns the extracted text as a string.
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
self._url("/api/v1/extract-pdf"),
|
||||
files={"file": ("document.pdf", pdf_bytes, "application/pdf")},
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return data.get("text", "")
|
||||
|
||||
|
||||
# Singleton
|
||||
embedding_client = EmbeddingClient()
|
||||
Reference in New Issue
Block a user