- RAG & SPIN Specialist improvements

eveai_chat_workers/retrievers/retriever_typing.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 
 from flask import current_app
 from pydantic import BaseModel, Field, model_validator
@@ -11,6 +11,7 @@ class RetrieverMetadata(BaseModel):
     document_id: int = Field(..., description="ID of the source document")
     version_id: int = Field(..., description="Version ID of the source document")
     document_name: str = Field(..., description="Name of the source document")
+    url: Optional[str] = Field(..., description="URL of the source document")
     user_metadata: Dict[str, Any] = Field(
         default_factory=dict,  # This will use an empty dict if None is provided
         description="User-defined metadata"

eveai_chat_workers/retrievers/standard_rag.py

@@ -97,6 +97,7 @@ class StandardRAGRetriever(BaseRetriever):
             query_obj = (
                 db.session.query(
                     db_class,
+                    DocumentVersion.url,
                     (1 - db_class.embedding.cosine_distance(query_embedding)).label('similarity')
                 )
                 .join(DocumentVersion, db_class.doc_vers_id == DocumentVersion.id)
@@ -116,7 +117,7 @@ class StandardRAGRetriever(BaseRetriever):
 
             # Transform results into standard format
             processed_results = []
-            for doc, similarity in results:
+            for doc, url, similarity in results:
                 # Parse user_metadata to ensure it's a dictionary
                 user_metadata = self._parse_metadata(doc.document_version.user_metadata)
                 processed_results.append(
@@ -128,6 +129,7 @@ class StandardRAGRetriever(BaseRetriever):
                             document_id=doc.document_version.doc_id,
                             version_id=doc.document_version.id,
                             document_name=doc.document_version.document.name,
+                            url=url or "",
                             user_metadata=user_metadata,
                         )
                     )