:sparkles: Add support for multi-round conversation

.gitignore

@@ -10,4 +10,5 @@ __pycache__
 uv.lock
 
 .env
-config.debug.yaml
+config.debug.yaml
+data/

app/models/__init__.py

@@ -1,9 +1 @@
-from .models import (
-    ChatCompletionRequest,
-    ChatCompletionResponse,
-    Message,
-    ModelData,
-    ModelListResponse,
-)
-
-__all__ = ["ChatCompletionRequest", "ChatCompletionResponse", "Message", "ModelData", "ModelListResponse"]
+from .models import *  # noqa: F403

app/models/models.py

@@ -1,6 +1,7 @@
+from datetime import datetime
 from typing import Dict, List, Literal, Optional, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class ContentItem(BaseModel):
@@ -20,21 +21,6 @@ class Message(BaseModel):
     name: Optional[str] = None
 
 
-class ChatCompletionRequest(BaseModel):
-    """Chat completion request model"""
-
-    model: str
-    messages: List[Message]
-    temperature: Optional[float] = 0.7
-    top_p: Optional[float] = 1.0
-    n: Optional[int] = 1
-    stream: Optional[bool] = False
-    max_tokens: Optional[int] = None
-    presence_penalty: Optional[float] = 0
-    frequency_penalty: Optional[float] = 0
-    user: Optional[str] = None
-
-
 class Choice(BaseModel):
     """Choice model"""
 
@@ -51,24 +37,39 @@ class Usage(BaseModel):
     total_tokens: int
 
 
-class ChatCompletionResponse(BaseModel):
-    """Chat completion response model"""
+class ModelData(BaseModel):
+    """Model data model"""
 
     id: str
-    object: str = "chat.completion"
+    object: str = "model"
     created: int
+    owned_by: str = "google"
+
+
+class ChatCompletionRequest(BaseModel):
+    """Chat completion request model"""
+
     model: str
-    choices: List[Choice]
-    usage: Usage
+    messages: List[Message]
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 1.0
+    n: Optional[int] = 1
+    stream: Optional[bool] = False
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = 0
+    frequency_penalty: Optional[float] = 0
+    user: Optional[str] = None
 
 
-class ModelData(BaseModel):
-    """Model data model"""
+class ChatCompletionResponse(BaseModel):
+    """Chat completion response model"""
 
     id: str
-    object: str = "model"
+    object: str = "chat.completion"
     created: int
-    owned_by: str = "google"
+    model: str
+    choices: List[Choice]
+    usage: Usage
 
 
 class ModelListResponse(BaseModel):
@@ -78,7 +79,21 @@ class ModelListResponse(BaseModel):
     data: List[ModelData]
 
 
-class ErrorResponse(BaseModel):
-    """Error response model"""
+class HealthCheckResponse(BaseModel):
+    """Health check response model"""
+
+    ok: bool
+    storage: Optional[Dict[str, str]] = None
+    error: Optional[str] = None
+
+
+class ConversationInStore(BaseModel):
+    """Conversation model for storing in the database."""
+
+    created_at: Optional[datetime] = Field(default=None)
+    updated_at: Optional[datetime] = Field(default=None)
 
-    error: Dict[str, str]
+    metadata: list[str | None] = Field(
+        ..., description="Metadata for Gemini API to locate the conversation"
+    )
+    messages: list[Message] = Field(..., description="Message contents in the conversation")

app/server/chat.py

@@ -1,19 +1,22 @@
-import json
-import time
 import uuid
 from datetime import datetime, timezone
 from pathlib import Path
 
+import orjson
 from fastapi import APIRouter, Depends, HTTPException, status
 from fastapi.responses import StreamingResponse
 from gemini_webapi.constants import Model
 from loguru import logger
 
-from ..models import ChatCompletionRequest, ModelData, ModelListResponse
-from ..services.client import SingletonGeminiClient
-from ..utils.utils import (
-    estimate_tokens,
+from ..models import (
+    ChatCompletionRequest,
+    ConversationInStore,
+    Message,
+    ModelData,
+    ModelListResponse,
 )
+from ..services import LMDBConversationStore, SingletonGeminiClient
+from ..utils.helper import estimate_tokens
 from .middleware import get_temp_dir, verify_api_key
 
 router = APIRouter()
@@ -47,46 +50,99 @@ async def create_chat_completion(
     tmp_dir: Path = Depends(get_temp_dir),
 ):
     client = SingletonGeminiClient()
+    db = LMDBConversationStore()
     model = Model.from_name(request.model)
 
-    # Preprocess the messages
-    try:
-        conversation, files = await client.prepare(request.messages, tmp_dir)
-        conversation = "\n".join(conversation)
-        logger.debug(f"Conversation length: {len(conversation)}, files count: {len(files)}")
-    except ValueError as e:
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
-    except Exception as e:
-        logger.exception(f"Error in preparing conversation: {e}")
-        raise
+    if len(request.messages) == 0:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="At least one message is required in the conversation.",
+        )
+
+    # Check if conversation is reusable
+    session = None
+    if _check_reusable(request.messages):
+        try:
+            # Exclude the last message from user
+            if old_conv := db.find(request.messages[:-1]):
+                session = client.start_chat(metadata=old_conv.metadata, model=model)
+        except Exception as e:
+            session = None
+            logger.warning(f"Error checking LMDB for reusable session: {e}")
+
+    if session:
+        # Just send the last message to the existing session
+        model_input, files = await client.process_message(
+            request.messages[-1], tmp_dir, tagged=False
+        )
+        logger.debug(f"Found reusable session: {session.metadata}")
+    else:
+        # Start a new session and concat messages into a single string
+        session = client.start_chat(model=model)
+        try:
+            model_input, files = await client.precess_conversation(request.messages, tmp_dir)
+        except ValueError as e:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
+        except Exception as e:
+            logger.exception(f"Error in preparing conversation: {e}")
+            raise
+        logger.debug("New session started.")
 
     # Generate response
     try:
-        response = await client.generate_content(conversation, files=files, model=model)
+        logger.debug(f"Input length: {len(model_input)}, files count: {len(files)}")
+        response = await session.send_message(model_input, files=files)
     except Exception as e:
         logger.exception(f"Error generating content from Gemini API: {e}")
         raise
 
-    # Post process
-    response_text = client.format_response(response)
-    if not response_text or response_text.strip() == "":
-        logger.warning("Empty response received from Gemini")
-        response_text = "No response generated."
+    # Format and clean the output
+    model_output = client.extract_output(response)
 
-    completion_id = f"chatcmpl-{uuid.uuid4()}"
-    timestamp = int(time.time())
+    # After cleaning, persist the conversation
+    try:
+        last_message = Message(role="assistant", content=model_output)
+        conv = ConversationInStore(
+            metadata=session.metadata,
+            messages=[*request.messages, last_message],
+        )
+        key = db.store(conv)
+        logger.debug(f"Conversation saved to LMDB with key: {key}")
+    except Exception as e:
+        # We can still return the response even if saving fails
+        logger.warning(f"Failed to save conversation to LMDB: {e}")
 
     # Return with streaming or standard response
+    completion_id = f"chatcmpl-{uuid.uuid4()}"
+    timestamp = int(datetime.now(tz=timezone.utc).timestamp())
     if request.stream:
-        return _create_streaming_response(response_text, completion_id, timestamp, request.model)
+        return _create_streaming_response(model_output, completion_id, timestamp, request.model)
     else:
         return _create_standard_response(
-            response_text, completion_id, timestamp, request.model, conversation
+            model_output, completion_id, timestamp, request.model, model_input
         )
 
 
+def _check_reusable(messages: list[Message]) -> bool:
+    """
+    Check if the conversation is reusable based on the message history.
+    """
+    if not messages or len(messages) < 2:
+        return False
+
+    # Last message must from the user
+    if messages[-1].role != "user" or not messages[-1].content:
+        return False
+
+    # The second last message must be from the assistant or system
+    if messages[-2].role not in ["assistant", "system"]:
+        return False
+
+    return True
+
+
 def _create_streaming_response(
-    response_text: str, completion_id: str, created_time: int, model: str
+    model_output: str, completion_id: str, created_time: int, model: str
 ) -> StreamingResponse:
     """Create streaming response"""
 
@@ -99,18 +155,18 @@ def _create_streaming_response(
             "model": model,
             "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
         }
-        yield f"data: {json.dumps(data)}\n\n"
+        yield f"data: {orjson.dumps(data).decode('utf-8')}\n\n"
 
         # Stream output text
-        for char in response_text:
+        for part in model_output.split():
             data = {
                 "id": completion_id,
                 "object": "chat.completion.chunk",
                 "created": created_time,
                 "model": model,
-                "choices": [{"index": 0, "delta": {"content": char}, "finish_reason": None}],
+                "choices": [{"index": 0, "delta": {"content": part}, "finish_reason": None}],
             }
-            yield f"data: {json.dumps(data)}\n\n"
+            yield f"data: {orjson.dumps(data).decode('utf-8')}\n\n"
 
         # Send end event
         data = {
@@ -120,19 +176,19 @@ def _create_streaming_response(
             "model": model,
             "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
         }
-        yield f"data: {json.dumps(data)}\n\n"
+        yield f"data: {orjson.dumps(data).decode('utf-8')}\n\n"
         yield "data: [DONE]\n\n"
 
     return StreamingResponse(generate_stream(), media_type="text/event-stream")
 
 
 def _create_standard_response(
-    response_text: str, completion_id: str, created_time: int, model: str, conversation: str
+    model_output: str, completion_id: str, created_time: int, model: str, model_input: str
 ) -> dict:
     """Create standard response"""
     # Calculate token usage
-    prompt_tokens = estimate_tokens(conversation)
-    completion_tokens = estimate_tokens(response_text)
+    prompt_tokens = estimate_tokens(model_input)
+    completion_tokens = estimate_tokens(model_output)
     total_tokens = prompt_tokens + completion_tokens
 
     result = {
@@ -143,7 +199,7 @@ def _create_standard_response(
         "choices": [
             {
                 "index": 0,
-                "message": {"role": "assistant", "content": response_text},
+                "message": {"role": "assistant", "content": model_output},
                 "finish_reason": "stop",
             }
         ],

app/server/health.py

@@ -1,7 +1,8 @@
 from fastapi import APIRouter
 from loguru import logger
 
-from ..services.client import SingletonGeminiClient
+from ..models import HealthCheckResponse
+from ..services import LMDBConversationStore, SingletonGeminiClient
 
 router = APIRouter()
 
@@ -9,15 +10,18 @@ router = APIRouter()
 @router.get("/health")
 async def health_check():
     client = SingletonGeminiClient()
+    db = LMDBConversationStore()
 
     if not client.running:
         try:
             await client.init()
         except Exception as e:
             logger.error(f"Failed to initialize Gemini client: {e}")
-            return {
-                "status": "unhealthy",
-                "error": str(e),
-            }
+            return HealthCheckResponse(ok=False, error=str(e))
 
-    return {"status": "healthy"}
+    stat = db.stats()
+    if not stat:
+        logger.error("Failed to retrieve LMDB conversation store stats")
+        return HealthCheckResponse(ok=False, error="LMDB conversation store unavailable")
+
+    return HealthCheckResponse(ok=True, storage=stat)

app/server/middleware.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 from fastapi import Depends, FastAPI, HTTPException, Request, status
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from fastapi.responses import ORJSONResponse
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 from ..utils import g_config
@@ -11,12 +11,12 @@ from ..utils import g_config
 
 def global_exception_handler(request: Request, exc: Exception):
     if isinstance(exc, HTTPException):
-        return JSONResponse(
+        return ORJSONResponse(
             status_code=exc.status_code,
             content={"error": {"message": exc.detail}},
         )
 
-    return JSONResponse(
+    return ORJSONResponse(
         status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content={"error": {"message": str(exc)}}
     )
 

app/services/__init__.py

@@ -0,0 +1,7 @@
+from .client import SingletonGeminiClient
+from .lmdb import LMDBConversationStore
+
+__all__ = [
+    "LMDBConversationStore",
+    "SingletonGeminiClient",
+]

app/services/client.py

@@ -5,63 +5,92 @@ from gemini_webapi import GeminiClient, ModelOutput
 
 from ..models import Message
 from ..utils import g_config
+from ..utils.helper import add_tag, save_file_to_tempfile, save_url_to_tempfile
 from ..utils.singleton import Singleton
-from ..utils.utils import add_tag, save_file_to_tempfile, save_url_to_tempfile
 
 
 class SingletonGeminiClient(GeminiClient, metaclass=Singleton):
     def __init__(self, **kwargs):
-        # TODO: Add proxy support if needed
-        super().__init__(
-            secure_1psid=g_config.gemini.secure_1psid,
-            secure_1psidts=g_config.gemini.secure_1psidts,
-            **kwargs,
-        )
-
-    async def init(self):
-        await super().init(
-            timeout=g_config.gemini.timeout,
-            auto_refresh=g_config.gemini.auto_refresh,
-            verbose=g_config.gemini.verbose,
-        )
-
-    async def prepare(self, messages: list[Message], tempdir: Path | None = None):
+        kwargs.setdefault("secure_1psid", g_config.gemini.secure_1psid)
+        kwargs.setdefault("secure_1psidts", g_config.gemini.secure_1psidts)
+
+        super().__init__(**kwargs)
+
+    async def init(self, **kwargs):
+        # Inject default configuration values
+        kwargs.setdefault("timeout", g_config.gemini.timeout)
+        kwargs.setdefault("auto_refresh", g_config.gemini.auto_refresh)
+        kwargs.setdefault("verbose", g_config.gemini.verbose)
+        kwargs.setdefault("refresh_interval", g_config.gemini.refresh_interval)
+
+        await super().init(**kwargs)
+
+    @staticmethod
+    async def process_message(
+        message: Message, tempdir: Path | None = None, tagged: bool = True
+    ) -> tuple[str, list[Path | str]]:
+        """
+        Process a single message and return model input.
+        """
+        model_input = ""
+        files: list[Path | str] = []
+        if isinstance(message.content, str):
+            # Pure text content
+            model_input = message.content
+        else:
+            # Mixed content
+            # TODO: Use Pydantic to enforce the value checking
+            for item in message.content:
+                if item.type == "text":
+                    model_input = item.text or ""
+
+                elif item.type == "image_url":
+                    if not item.image_url:
+                        raise ValueError("Image URL cannot be empty")
+                    if url := item.image_url.get("url", None):
+                        files.append(await save_url_to_tempfile(url, tempdir))
+                    else:
+                        raise ValueError("Image URL must contain 'url' key")
+
+                elif item.type == "file":
+                    if not item.file:
+                        raise ValueError("File cannot be empty")
+                    if file_data := item.file.get("file_data", None):
+                        filename = item.file.get("filename", "")
+                        files.append(await save_file_to_tempfile(file_data, filename, tempdir))
+                    else:
+                        raise ValueError("File must contain 'file_data' key")
+
+        # Add role tag if needed
+        if model_input and tagged:
+            model_input = add_tag(message.role, model_input)
+
+        return model_input, files
+
+    @staticmethod
+    async def precess_conversation(messages: list[Message], tempdir: Path | None = None):
+        """
+        Process the entire conversation and return a formatted string and list of files.
+        The last message is assumed to be the assistant's response.
+        """
         conversation: list[str] = []
         files: list[Path | str] = []
 
         for msg in messages:
-            if isinstance(msg.content, str):
-                # Pure text content
-                conversation.append(add_tag(msg.role, msg.content))
-            else:
-                # Mixed content
-                for item in msg.content:
-                    if item.type == "text":
-                        conversation.append(add_tag(msg.role, item.text or ""))
-
-                    elif item.type == "image_url":
-                        # TODO: Use Pydantic to enforce the value checking
-                        if not item.image_url:
-                            raise ValueError("Image URL cannot be empty")
-                        if url := item.image_url.get("url", None):
-                            files.append(await save_url_to_tempfile(url, tempdir))
-                        else:
-                            raise ValueError("Image URL must contain 'url' key")
-
-                    elif item.type == "file":
-                        if not item.file:
-                            raise ValueError("File cannot be empty")
-                        if file_data := item.file.get("file_data", None):
-                            filename = item.file.get("filename", "")
-                            files.append(await save_file_to_tempfile(file_data, filename, tempdir))
-                        else:
-                            raise ValueError("File must contain 'file_data' key")
+            input_part, files_part = await SingletonGeminiClient.process_message(msg, tempdir)
+            conversation.append(input_part)
+            files.extend(files_part)
 
         # Left with the last message as the assistant's response
-        conversation.append(add_tag("assistant", "", open=True))
-        return conversation, files
+        conversation.append(add_tag("assistant", "", unclose=True))
+
+        return "\n".join(conversation), files
 
-    def format_response(self, response: ModelOutput):
+    @staticmethod
+    def extract_output(response: ModelOutput):
+        """
+        Extract and format the output text from the Gemini response.
+        """
         text = ""
 
         if response.thoughts:
@@ -76,14 +105,12 @@ class SingletonGeminiClient(GeminiClient, metaclass=Singleton):
         text = text.replace("&lt;", "<").replace("\\<", "<").replace("\\_", "_").replace("\\>", ">")
 
         def simplify_link_target(text_content: str) -> str:
-            """简化链接目标"""
             match_colon_num = re.match(r"([^:]+:\d+)", text_content)
             if match_colon_num:
                 return match_colon_num.group(1)
             return text_content
 
         def replacer(match: re.Match) -> str:
-            """链接替换器"""
             outer_open_paren = match.group(1)
             display_text = match.group(2)
 
@@ -95,10 +122,10 @@ class SingletonGeminiClient(GeminiClient, metaclass=Singleton):
             else:
                 return new_link_segment
 
-        # 修复Google搜索链接
+        # Replace Google search links with simplified markdown links
         pattern = r"(\()?\[`([^`]+?)`\]\((https://www.google.com/search\?q=)(.*?)(?<!\\)\)\)*(\))?"
         text = re.sub(pattern, replacer, text)
 
-        # 修复包装的markdown链接
+        # Fix inline code blocks
         pattern = r"`(\[[^\]]+\]\([^\)]+\))`"
         return re.sub(pattern, r"\1", text)

app/services/lmdb.py

@@ -0,0 +1,325 @@
+import hashlib
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import lmdb
+import orjson
+from loguru import logger
+
+from ..models import ConversationInStore, Message
+from ..utils import g_config
+from ..utils.singleton import Singleton
+
+
+def hash_message(message: Message) -> str:
+    """Generate a hash for a single message."""
+    # Convert message to dict and sort keys for consistent hashing
+    message_dict = message.model_dump(mode="json")
+    message_bytes = orjson.dumps(message_dict, option=orjson.OPT_SORT_KEYS)
+    return hashlib.sha256(message_bytes).hexdigest()
+
+
+def hash_message_list(messages: List[Message]) -> str:
+    """Generate a hash for a list of messages."""
+    # Create a combined hash from all individual message hashes
+    combined_hash = hashlib.sha256()
+    for message in messages:
+        message_hash = hash_message(message)
+        combined_hash.update(message_hash.encode("utf-8"))
+    return combined_hash.hexdigest()
+
+
+class LMDBConversationStore(metaclass=Singleton):
+    """LMDB-based storage for Message lists with hash-based key-value operations."""
+
+    CUSTOM_KEY_BINDING_PREFIX = "hash:"
+
+    def __init__(self, db_path: Optional[str] = None, max_db_size: Optional[int] = None):
+        """
+        Initialize LMDB store.
+
+        Args:
+            db_path: Path to LMDB database directory
+            max_db_size: Maximum database size in bytes (default: 128MB)
+        """
+
+        if db_path is None:
+            db_path = g_config.storage.path
+        if max_db_size is None:
+            max_db_size = g_config.storage.max_size
+
+        self.db_path: Path = Path(db_path)
+        self.max_db_size: int = max_db_size
+        self._env: lmdb.Environment | None = None
+
+        self._ensure_db_path()
+        self._init_environment()
+
+    def _ensure_db_path(self) -> None:
+        """Ensure database directory exists."""
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+
+    def _init_environment(self) -> None:
+        """Initialize LMDB environment."""
+        try:
+            self._env = lmdb.open(
+                str(self.db_path),
+                map_size=self.max_db_size,
+                max_dbs=3,  # main, metadata, and index databases
+                writemap=True,
+                readahead=False,
+                meminit=False,
+            )
+            logger.info(f"LMDB environment initialized at {self.db_path}")
+        except Exception as e:
+            logger.error(f"Failed to initialize LMDB environment: {e}")
+            raise
+
+    @contextmanager
+    def _get_transaction(self, write: bool = False):
+        """Get LMDB transaction context manager."""
+        if not self._env:
+            raise RuntimeError("LMDB environment not initialized")
+
+        txn: lmdb.Transaction = self._env.begin(write=write)
+        try:
+            yield txn
+            if write:
+                txn.commit()
+        except Exception:
+            if write:
+                txn.abort()
+            raise
+        finally:
+            pass  # Transaction is automatically cleaned up
+
+    def store(
+        self,
+        conv: ConversationInStore,
+        custom_key: Optional[str] = None,
+    ) -> str:
+        """
+        Store a conversation model in LMDB.
+
+        Args:
+            conv: Conversation model to store
+            custom_key: Optional custom key, if not provided, hash will be used
+
+        Returns:
+            str: The key used to store the messages (hash or custom key)
+        """
+        if not conv:
+            raise ValueError("Messages list cannot be empty")
+
+        # Generate hash for the message list
+        message_hash = hash_message_list(conv.messages)
+        storage_key = custom_key or message_hash
+
+        # Prepare data for storage
+        now = datetime.now()
+        if conv.created_at is None:
+            conv.created_at = now
+        conv.updated_at = now
+
+        value = orjson.dumps(conv.model_dump(mode="json"))
+
+        try:
+            with self._get_transaction(write=True) as txn:
+                # Store main data
+                txn.put(storage_key.encode("utf-8"), value, overwrite=True)
+
+                # Store hash -> key mapping for reverse lookup
+                if custom_key:
+                    txn.put(
+                        f"{self.CUSTOM_KEY_BINDING_PREFIX}{message_hash}".encode("utf-8"),
+                        custom_key.encode("utf-8"),
+                    )
+
+                logger.debug(f"Stored {len(conv.messages)} messages with key: {storage_key}")
+                return storage_key
+
+        except Exception as e:
+            logger.error(f"Failed to store conversation: {e}")
+            raise
+
+    def get(self, key: str) -> Optional[ConversationInStore]:
+        """
+        Retrieve conversation data by key.
+
+        Args:
+            key: Storage key (hash or custom key)
+
+        Returns:
+            Conversation or None if not found
+        """
+        try:
+            with self._get_transaction(write=False) as txn:
+                data = txn.get(key.encode("utf-8"), default=None)
+                if not data:
+                    return None
+
+                storage_data = orjson.loads(data)  # type: ignore
+                conv = ConversationInStore.model_validate(storage_data)
+
+                logger.debug(f"Retrieved {len(conv.messages)} messages for key: {key}")
+                return conv
+
+        except Exception as e:
+            logger.error(f"Failed to retrieve messages for key {key}: {e}")
+            return None
+
+    def find(self, messages: List[Message]) -> Optional[ConversationInStore]:
+        """
+        Search conversation data by message list.
+
+        Args:
+            messages: List of messages to search for
+
+        Returns:
+            Conversation or None if not found
+        """
+        if not messages:
+            return None
+
+        message_hash = hash_message_list(messages)
+        key = f"{self.CUSTOM_KEY_BINDING_PREFIX}{message_hash}"
+
+        try:
+            with self._get_transaction(write=False) as txn:
+                # Try custom key binding first
+                key = txn.get(key.encode("utf-8"), default=None)
+                key = key.decode("utf-8") if key else message_hash  # type: ignore
+
+                # Fallback to hash if no custom key found
+                return self.get(key)
+
+        except Exception as e:
+            logger.error(f"Failed to retrieve messages by message list: {e}")
+            return None
+
+    def exists(self, key: str) -> bool:
+        """
+        Check if a key exists in the store.
+
+        Args:
+            key: Storage key to check
+
+        Returns:
+            bool: True if key exists, False otherwise
+        """
+        try:
+            with self._get_transaction(write=False) as txn:
+                return txn.get(key.encode("utf-8")) is not None
+        except Exception as e:
+            logger.error(f"Failed to check existence of key {key}: {e}")
+            return False
+
+    def delete(self, key: str) -> Optional[ConversationInStore]:
+        """
+        Delete conversation model by key.
+
+        Args:
+            key: Storage key to delete
+
+        Returns:
+            ConversationInStore: The deleted conversation data, or None if not found
+        """
+        try:
+            with self._get_transaction(write=True) as txn:
+                # Get data first to clean up hash mapping
+                data = txn.get(key.encode("utf-8"))
+                if not data:
+                    return None
+
+                storage_data = orjson.loads(data)  # type: ignore
+                conv = ConversationInStore.model_validate(storage_data)
+                message_hash = hash_message_list(conv.messages)
+
+                # Delete main data
+                txn.delete(key.encode("utf-8"))
+
+                # Clean up hash mapping if it exists
+                if message_hash and key != message_hash:
+                    txn.delete(f"{self.CUSTOM_KEY_BINDING_PREFIX}{message_hash}".encode("utf-8"))
+
+                logger.info(f"Deleted messages with key: {key}")
+                return conv
+
+        except Exception as e:
+            logger.error(f"Failed to delete key {key}: {e}")
+            return None
+
+    def keys(self, prefix: str = "", limit: Optional[int] = None) -> List[str]:
+        """
+        List all keys in the store, optionally filtered by prefix.
+
+        Args:
+            prefix: Optional prefix to filter keys
+            limit: Optional limit on number of keys returned
+
+        Returns:
+            List of keys
+        """
+        keys = []
+        try:
+            with self._get_transaction(write=False) as txn:
+                cursor = txn.cursor()
+                cursor.first()
+
+                count = 0
+                for key, _ in cursor:
+                    key_str = key.decode("utf-8")
+                    # Skip internal hash mappings
+                    if key_str.startswith(self.CUSTOM_KEY_BINDING_PREFIX):
+                        continue
+
+                    if not prefix or key_str.startswith(prefix):
+                        keys.append(key_str)
+                        count += 1
+
+                        if limit and count >= limit:
+                            break
+
+        except Exception as e:
+            logger.error(f"Failed to list keys: {e}")
+
+        return keys
+
+    def stats(self) -> Dict[str, Any]:
+        """
+        Get database statistics.
+
+        Returns:
+            Dict with database statistics
+        """
+        if not self._env:
+            logger.error("LMDB environment not initialized")
+            return {}
+
+        try:
+            with self._get_transaction(write=False) as txn:
+                stat = txn.stat(self._env._db)
+                return {
+                    "entries": stat["entries"],
+                    "page_size": stat["psize"],
+                    "depth": stat["depth"],
+                    "branch_pages": stat["branch_pages"],
+                    "leaf_pages": stat["leaf_pages"],
+                    "overflow_pages": stat["overflow_pages"],
+                }
+        except Exception as e:
+            logger.error(f"Failed to get database stats: {e}")
+            return {}
+
+    def close(self) -> None:
+        """Close the LMDB environment."""
+        if self._env:
+            self._env.close()
+            self._env = None
+            logger.info("LMDB environment closed")
+
+    def __del__(self):
+        """Cleanup on destruction."""
+        self.close()

app/utils/config.py

@@ -23,10 +23,13 @@ class ServerConfig(BaseModel):
 class GeminiConfig(BaseModel):
     """Gemini API configuration"""
 
-    secure_1psid: str = Field(..., description="Gemini API Secure 1PSID")
-    secure_1psidts: str = Field(..., description="Gemini API Secure 1PSIDTS")
+    secure_1psid: str = Field(..., description="Gemini Secure 1PSID")
+    secure_1psidts: str = Field(..., description="Gemini Secure 1PSIDTS")
     timeout: int = Field(default=60, ge=1, description="Init timeout")
-    auto_refresh: bool = Field(True, description="Enable auto-refresh for Gemini API credentials")
+    auto_refresh: bool = Field(True, description="Enable auto-refresh for Gemini cookies")
+    refresh_interval: int = Field(
+        default=540, ge=1, description="Interval in seconds to refresh Gemini cookies"
+    )
     verbose: bool = Field(False, description="Enable verbose logging for Gemini API requests")
 
 
@@ -45,6 +48,18 @@ class CORSConfig(BaseModel):
     )
 
 
+class StorageConfig(BaseModel):
+    path: str = Field(
+        default="data/msg.lmdb",
+        description="Path to the storage directory where data will be saved",
+    )
+    max_size: int = Field(
+        default=1024**2 * 128,  # 128 MB
+        ge=1,
+        description="Maximum size of the storage in bytes",
+    )
+
+
 class LoggingConfig(BaseModel):
     level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = Field(
         default="DEBUG",
@@ -70,6 +85,11 @@ class Config(BaseSettings):
     # Gemini API configuration
     gemini: GeminiConfig = Field(..., description="Gemini API configuration, must be set")
 
+    storage: StorageConfig = Field(
+        default=StorageConfig(),
+        description="Storage configuration, defines where and how data will be stored",
+    )
+
     # Logging configuration
     logging: LoggingConfig = Field(
         default=LoggingConfig(),

app/utils/{utils.py → helper.py}

@@ -6,13 +6,13 @@ import httpx
 from loguru import logger
 
 
-def add_tag(role: str, content: str, open: bool = False) -> str:
+def add_tag(role: str, content: str, unclose: bool = False) -> str:
     """Surround content with role tags"""
     if role not in ["user", "assistant", "system"]:
         logger.warning(f"Unknown role: {role}, returning content without tags")
         return content
 
-    return f"<|im_start|>{role}\n{content}" + ("\n<|im_end|>" if not open else "")
+    return f"<|im_start|>{role}\n{content}" + ("\n<|im_end|>" if not unclose else "")
 
 
 def estimate_tokens(text: str) -> int:

app/utils/logging.py

@@ -3,7 +3,6 @@ import logging
 import sys
 from typing import Literal
 
-from gemini_webapi.utils.logger import set_log_level
 from loguru import logger
 
 
@@ -22,10 +21,7 @@ def setup_logging(
         backtrace: Whether to enable backtrace information
         colorize: Whether to enable colors
     """
-    # Set gemini_webapi log level to avoid conflicts
-    set_log_level(level)
-
-    # Remove all existing loguru handlers
+    # Reset all logger handlers
     logger.remove()
 
     # Add unified handler for all logs
@@ -69,4 +65,4 @@ def _setup_logging_intercept() -> None:
             logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage())
 
     # Remove all existing handlers and add our interceptor
-    logging.basicConfig(handlers=[InterceptHandler()], level=0, force=True)
+    logging.basicConfig(handlers=[InterceptHandler()], level="INFO", force=True)

config/config.yaml

@@ -11,12 +11,17 @@ cors:
   allow_methods: ["*"]
   allow_headers: ["*"]
 
+storage:
+  path: "data/lmdb"
+  max_size: 134217728 # 128 MB
+
 gemini:
   secure_1psid: "YOUR_SECURE_1PSID_HERE"
   secure_1psidts: "YOUR_SECURE_1PSIDTS_HERE"
   timeout: 60
   auto_refresh: true
+  refresh_interval: 540
   verbose: false
 
 logging:
-  level: "DEBUG"
+  level: "DEBUG"

pyproject.toml

@@ -6,9 +6,10 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "fastapi>=0.115.12",
-    "gemini-webapi>=1.12.1",
+    "gemini-webapi>=1.14.0",
+    "lmdb>=1.6.2",
     "loguru>=0.7.0",
-    "pydantic-settings>=2.9.1",
+    "pydantic-settings[yaml]>=2.9.1",
     "uvicorn>=0.34.1",
     "uvloop>=0.21.0",
 ]