94 lines
2.6 KiB
Python
94 lines
2.6 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="Current conversation messages (user and assistant only)",
|
|
)
|
|
notifications: list[str] | None = Field(
|
|
default=None,
|
|
description="Active notifications for the user",
|
|
)
|
|
conversation_history: str | None = Field(
|
|
default=None,
|
|
description="Formatted conversation history",
|
|
)
|
|
user_nickname: str | None = Field(
|
|
default=None,
|
|
description="User's nickname or display name",
|
|
)
|
|
|
|
|
|
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]],
|
|
notifications: list[str] | None = None,
|
|
conversation_history: str | None = None,
|
|
user_nickname: str | None = None,
|
|
) -> str:
|
|
"""Send conversation to RAG endpoint and get response.
|
|
|
|
Args:
|
|
messages: Current conversation messages (user/assistant only)
|
|
e.g., [{"role": "user", "content": "Hello"}, ...]
|
|
notifications: Active notifications for the user (optional)
|
|
conversation_history: Formatted conversation history (optional)
|
|
user_nickname: User's nickname or display name (optional)
|
|
|
|
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()
|