70 lines
1.8 KiB
Python
70 lines
1.8 KiB
Python
"""Base RAG service interface."""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from types import TracebackType
|
|
from typing import Self
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class Message(BaseModel):
|
|
"""OpenAI-style message format."""
|
|
|
|
role: str = Field(..., description="Role: system, user, or assistant")
|
|
content: str = Field(..., description="Message content")
|
|
|
|
|
|
class RAGRequest(BaseModel):
|
|
"""Request model for RAG endpoint."""
|
|
|
|
messages: list[Message] = Field(..., description="Conversation history")
|
|
|
|
|
|
class RAGResponse(BaseModel):
|
|
"""Response model from RAG endpoint."""
|
|
|
|
response: str = Field(..., description="Generated response from RAG")
|
|
|
|
|
|
class RAGServiceBase(ABC):
|
|
"""Abstract base class for RAG service implementations.
|
|
|
|
Provides a common interface for different RAG service backends
|
|
(HTTP, mock, echo, etc.).
|
|
"""
|
|
|
|
@abstractmethod
|
|
async def query(self, messages: list[dict[str, str]]) -> str:
|
|
"""Send conversation history to RAG endpoint and get response.
|
|
|
|
Args:
|
|
messages: OpenAI-style conversation history
|
|
e.g., [{"role": "user", "content": "Hello"}, ...]
|
|
|
|
Returns:
|
|
Response string from RAG endpoint
|
|
|
|
Raises:
|
|
Exception: Implementation-specific exceptions
|
|
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
async def close(self) -> None:
|
|
"""Close the service and release resources."""
|
|
...
|
|
|
|
async def __aenter__(self) -> Self:
|
|
"""Async context manager entry."""
|
|
return self
|
|
|
|
async def __aexit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_val: BaseException | None,
|
|
exc_tb: TracebackType | None,
|
|
) -> None:
|
|
"""Async context manager exit."""
|
|
await self.close()
|