Doc RAG updates

This commit is contained in:
Cyber MacGeddon 2026-03-07 15:01:52 +00:00
parent 01eef99218
commit 9f90075367

View file

@ -11,6 +11,11 @@ Document embeddings storage currently stores chunk text directly in the vector s
class ChunkEmbeddings:
chunk: bytes = b""
vectors: list[list[float]] = field(default_factory=list)
@dataclass
class DocumentEmbeddingsResponse:
error: Error | None = None
chunks: list[str] = field(default_factory=list)
```
Vector store payload:
@ -18,6 +23,10 @@ Vector store payload:
payload={"doc": chunk} # Duplicates Garage content
```
Document RAG flow:
1. Query embeddings → get chunk text
2. Pass chunk text directly to prompt
## Design
### Schema Changes
@ -45,20 +54,55 @@ All stores (Qdrant, Milvus, Pinecone):
payload={"chunk_id": chunk_id}
```
### Query Flow
### Document RAG Changes
1. Search vector store → get matching `chunk_id` values
2. Return chunk_ids to caller
3. Caller fetches content from Garage if needed
The document RAG processor must fetch chunk content from Garage:
```python
# In document_rag.py get_docs():
async def get_docs(self, query):
vectors = await self.get_vector(query)
# Get chunk_ids from embeddings store
chunk_ids = await self.rag.doc_embeddings_client.query(
vectors, limit=self.doc_limit,
user=self.user, collection=self.collection,
)
# Fetch chunk content from Garage
docs = []
for chunk_id in chunk_ids:
content = await self.rag.librarian_client.get_document_content(
chunk_id, self.user
)
docs.append(content)
return docs
```
### Client Changes
**DocumentEmbeddingsClient** - return chunk_ids:
```python
async def query(self, vectors, limit=20, user="trustgraph",
collection="default", timeout=30):
resp = await self.request(...)
if resp.error:
raise RuntimeError(resp.error.message)
return resp.chunk_ids # Changed from resp.chunks
```
## Files to Modify
### Schema
- `trustgraph-base/trustgraph/schema/knowledge/embeddings.py`
- `trustgraph-base/trustgraph/schema/services/query.py`
- `trustgraph-base/trustgraph/schema/knowledge/embeddings.py` - ChunkEmbeddings
- `trustgraph-base/trustgraph/schema/services/query.py` - DocumentEmbeddingsResponse
### Client
- `trustgraph-base/trustgraph/base/document_embeddings_client.py` - return chunk_ids
### Embeddings Service
- `trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py`
- `trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py` - pass chunk_id
### Storage Writers
- `trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py`
@ -70,6 +114,10 @@ payload={"chunk_id": chunk_id}
- `trustgraph-flow/trustgraph/query/doc_embeddings/milvus/service.py`
- `trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py`
### Document RAG
- `trustgraph-flow/trustgraph/retrieval/document_rag/rag.py` - add librarian client
- `trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py` - fetch from Garage
## Benefits
1. Single source of truth - chunk text only in Garage