Files
eveAI/eveai_chat_workers/retrievers/base.py
Josako 25213f2004 - Implementation of specialist execution api, including SSE protocol
- eveai_chat becomes deprecated and should be replaced with SSE
- Adaptation of STANDARD_RAG specialist
- Base class definition allowing to realise specialists with crewai framework
- Implementation of SPIN_SPECIALIST
- Implementation of test app for testing specialists (test_specialist_client). Also serves as an example for future SSE-based client
- Improvements to startup scripts to better handle and scale multiple connections
- Small improvements to the interaction forms and views
- Caching implementation improved and augmented with additional caches
2025-02-20 05:50:16 +01:00

58 lines
1.9 KiB
Python

from abc import ABC, abstractmethod, abstractproperty
from typing import Dict, Any, List
from flask import current_app
from eveai_chat_workers.retrievers.retriever_typing import RetrieverResult, RetrieverArguments
from config.logging_config import TuningLogger
class BaseRetriever(ABC):
"""Base class for all retrievers"""
def __init__(self, tenant_id: int, retriever_id: int):
self.tenant_id = tenant_id
self.retriever_id = retriever_id
self.tuning = False
self.tuning_logger = None
self._setup_tuning_logger()
@property
@abstractmethod
def type(self) -> str:
"""The type of the retriever"""
pass
def _setup_tuning_logger(self):
try:
self.tuning_logger = TuningLogger(
'tuning',
tenant_id=self.tenant_id,
retriever_id=self.retriever_id,
)
# Verify logger is working with a test message
if self.tuning:
self.tuning_logger.log_tuning('retriever', "Tuning logger initialized")
except Exception as e:
current_app.logger.error(f"Failed to setup tuning logger: {str(e)}")
raise
def log_tuning(self, message: str, data: Dict[str, Any] = None) -> None:
if self.tuning and self.tuning_logger:
try:
self.tuning_logger.log_tuning('retriever', message, data)
except Exception as e:
current_app.logger.error(f"Processor: Error in tuning logging: {e}")
@abstractmethod
def retrieve(self, arguments: RetrieverArguments) -> List[RetrieverResult]:
"""
Retrieve relevant documents based on provided arguments
Args:
arguments: Dictionary of arguments for the retrieval operation
Returns:
List[Dict[str, Any]]: List of retrieved documents/content
"""
pass