v0

backend/storage.py

@@ -0,0 +1,154 @@
+"""JSON-based storage for conversations."""
+
+import json
+import os
+from datetime import datetime
+from typing import List, Dict, Any, Optional
+from pathlib import Path
+from .config import DATA_DIR
+
+
+def ensure_data_dir():
+    """Ensure the data directory exists."""
+    Path(DATA_DIR).mkdir(parents=True, exist_ok=True)
+
+
+def get_conversation_path(conversation_id: str) -> str:
+    """Get the file path for a conversation."""
+    return os.path.join(DATA_DIR, f"{conversation_id}.json")
+
+
+def create_conversation(conversation_id: str) -> Dict[str, Any]:
+    """
+    Create a new conversation.
+
+    Args:
+        conversation_id: Unique identifier for the conversation
+
+    Returns:
+        New conversation dict
+    """
+    ensure_data_dir()
+
+    conversation = {
+        "id": conversation_id,
+        "created_at": datetime.utcnow().isoformat(),
+        "messages": []
+    }
+
+    # Save to file
+    path = get_conversation_path(conversation_id)
+    with open(path, 'w') as f:
+        json.dump(conversation, f, indent=2)
+
+    return conversation
+
+
+def get_conversation(conversation_id: str) -> Optional[Dict[str, Any]]:
+    """
+    Load a conversation from storage.
+
+    Args:
+        conversation_id: Unique identifier for the conversation
+
+    Returns:
+        Conversation dict or None if not found
+    """
+    path = get_conversation_path(conversation_id)
+
+    if not os.path.exists(path):
+        return None
+
+    with open(path, 'r') as f:
+        return json.load(f)
+
+
+def save_conversation(conversation: Dict[str, Any]):
+    """
+    Save a conversation to storage.
+
+    Args:
+        conversation: Conversation dict to save
+    """
+    ensure_data_dir()
+
+    path = get_conversation_path(conversation['id'])
+    with open(path, 'w') as f:
+        json.dump(conversation, f, indent=2)
+
+
+def list_conversations() -> List[Dict[str, Any]]:
+    """
+    List all conversations (metadata only).
+
+    Returns:
+        List of conversation metadata dicts
+    """
+    ensure_data_dir()
+
+    conversations = []
+    for filename in os.listdir(DATA_DIR):
+        if filename.endswith('.json'):
+            path = os.path.join(DATA_DIR, filename)
+            with open(path, 'r') as f:
+                data = json.load(f)
+                # Return metadata only
+                conversations.append({
+                    "id": data["id"],
+                    "created_at": data["created_at"],
+                    "message_count": len(data["messages"])
+                })
+
+    # Sort by creation time, newest first
+    conversations.sort(key=lambda x: x["created_at"], reverse=True)
+
+    return conversations
+
+
+def add_user_message(conversation_id: str, content: str):
+    """
+    Add a user message to a conversation.
+
+    Args:
+        conversation_id: Conversation identifier
+        content: User message content
+    """
+    conversation = get_conversation(conversation_id)
+    if conversation is None:
+        raise ValueError(f"Conversation {conversation_id} not found")
+
+    conversation["messages"].append({
+        "role": "user",
+        "content": content
+    })
+
+    save_conversation(conversation)
+
+
+def add_assistant_message(
+    conversation_id: str,
+    stage1: List[Dict[str, Any]],
+    stage2: List[Dict[str, Any]],
+    stage3: Dict[str, Any]
+):
+    """
+    Add an assistant message with all 3 stages to a conversation.
+
+    Args:
+        conversation_id: Conversation identifier
+        stage1: List of individual model responses
+        stage2: List of model rankings
+        stage3: Final synthesized response
+    """
+    conversation = get_conversation(conversation_id)
+    if conversation is None:
+        raise ValueError(f"Conversation {conversation_id} not found")
+
+    conversation["messages"].append({
+        "role": "assistant",
+        "stage1": stage1,
+        "stage2": stage2,
+        "stage3": stage3
+    })
+
+    save_conversation(conversation)