Update utils/case_manager.py

utils/case_manager.py

@@ -1,108 +1,99 @@
 import os
 import json
 from datetime import datetime
-from typing import Dict, List, Optional
-import streamlit as st
+from typing import List, Dict, Optional
 
 
 class CaseManager:
     def __init__(self, base_path: str = "data/cases"):
+        """
+        Initialize CaseManager with a base directory to store cases.
+        """
         self.base_path = base_path
         os.makedirs(base_path, exist_ok=True)
         self._load_cases()
 
     def _load_cases(self):
         """Load existing cases from storage."""
-        if "cases" not in st.session_state:
-            st.session_state.cases = {}
-            for case_id in os.listdir(self.base_path):
-                case_path = os.path.join(self.base_path, case_id)
-                if os.path.isdir(case_path):
-                    with open(os.path.join(case_path, "metadata.json"), "r") as f:
-                        st.session_state.cases[case_id] = json.load(f)
+        self.cases = {}
+        for case_id in os.listdir(self.base_path):
+            case_path = os.path.join(self.base_path, case_id)
+            if os.path.isdir(case_path):
+                with open(os.path.join(case_path, 'metadata.json'), 'r') as f:
+                    self.cases[case_id] = json.load(f)
 
     def create_case(self, title: str, description: str, case_type: str) -> str:
-        """Create a new case."""
+        """
+        Create a new case and save it to storage.
+        """
         case_id = datetime.now().strftime('%Y%m%d_%H%M%S')
         case_path = os.path.join(self.base_path, case_id)
         os.makedirs(case_path, exist_ok=True)
-        os.makedirs(os.path.join(case_path, "documents"), exist_ok=True)
+        os.makedirs(os.path.join(case_path, 'documents'), exist_ok=True)
 
         case_data = {
-            "id": case_id,
-            "title": title,
-            "description": description,
-            "case_type": case_type,
-            "created_at": datetime.now().isoformat(),
-            "documents": []
+            'id': case_id,
+            'title': title,
+            'description': description,
+            'case_type': case_type,
+            'created_at': datetime.now().isoformat(),
+            'documents': []
         }
 
-        with open(os.path.join(case_path, "metadata.json"), "w") as f:
+        with open(os.path.join(case_path, 'metadata.json'), 'w') as f:
             json.dump(case_data, f)
-
-        st.session_state.cases[case_id] = case_data
+        self.cases[case_id] = case_data
         return case_id
 
-    def add_document(self, case_id: str, document_data: Dict) -> str:
-        """Add a document to a case."""
-        case = st.session_state.cases.get(case_id)
+    def add_document(self, case_id: str, document_data: Dict):
+        """
+        Add a document to a case and save it to storage.
+        """
+        case = self.cases.get(case_id)
         if not case:
-            raise ValueError(f"Case {case_id} not found")
+            raise ValueError(f"Case with ID {case_id} not found.")
 
-        doc_id = datetime.now().strftime('%Y%m%d_%H%M%S')
-        document_data["id"] = doc_id
-        document_data["added_at"] = datetime.now().isoformat()
+        # Ensure the document has required fields
+        document_data['id'] = document_data.get('id', datetime.now().strftime('%Y%m%d_%H%M%S'))
+        document_data['added_at'] = document_data.get('added_at', datetime.now().isoformat())
+        document_data['text'] = document_data.get('text', '')  # Ensure text is always present
 
-        # Save document to case directory
-        case_path = os.path.join(self.base_path, case_id)
-        doc_path = os.path.join(case_path, "documents", f"{doc_id}.json")
-        with open(doc_path, "w") as f:
-            json.dump(document_data, f)
+        case['documents'].append(document_data)
 
-        # Update case metadata
-        case["documents"].append({"id": doc_id, "title": document_data["title"], "added_at": document_data["added_at"]})
-        with open(os.path.join(case_path, "metadata.json"), "w") as f:
+        case_path = os.path.join(self.base_path, case_id)
+        with open(os.path.join(case_path, 'metadata.json'), 'w') as f:
             json.dump(case, f)
 
-        return doc_id
+    def list_documents(self, case_id: str) -> List[Dict]:
+        """
+        List all documents for a given case.
+        """
+        case = self.cases.get(case_id)
+        if not case:
+            raise ValueError(f"Case with ID {case_id} not found.")
+        return case.get('documents', [])
 
     def get_case(self, case_id: str) -> Optional[Dict]:
-        """Get case details."""
-        return st.session_state.cases.get(case_id)
+        """
+        Get details of a specific case by ID.
+        """
+        return self.cases.get(case_id)
 
     def get_all_cases(self) -> List[Dict]:
-        """Get all cases."""
-        return list(st.session_state.cases.values())
-
-    def get_document(self, case_id: str, doc_id: str) -> Optional[Dict]:
-        """Retrieve a specific document by ID."""
-        case_path = os.path.join(self.base_path, case_id)
-        doc_path = os.path.join(case_path, "documents", f"{doc_id}.json")
-        if os.path.exists(doc_path):
-            with open(doc_path, "r") as f:
-                return json.load(f)
-        return None
-
-    def delete_document(self, case_id: str, doc_id: str):
-        """Delete a specific document."""
-        case = st.session_state.cases.get(case_id)
-        if not case:
-            raise ValueError(f"Case {case_id} not found")
-
-        # Remove document file
-        case_path = os.path.join(self.base_path, case_id)
-        doc_path = os.path.join(case_path, "documents", f"{doc_id}.json")
-        if os.path.exists(doc_path):
-            os.remove(doc_path)
-
-        # Update case metadata
-        case["documents"] = [doc for doc in case["documents"] if doc["id"] != doc_id]
-        with open(os.path.join(case_path, "metadata.json"), "w") as f:
-            json.dump(case, f)
-
-    def list_documents(self, case_id: str) -> List[Dict]:
-        """List all documents for a case."""
-        case = st.session_state.cases.get(case_id)
-        if not case:
-            return []
-        return case["documents"]
+        """
+        Get a list of all cases.
+        """
+        return list(self.cases.values())
+
+    def search(self, query: str) -> List[Dict]:
+        """
+        Search for cases or documents by title or type.
+        """
+        results = []
+        for case in self.cases.values():
+            if query.lower() in case['title'].lower() or query.lower() in case['case_type'].lower():
+                results.append(case)
+            for document in case['documents']:
+                if query.lower() in document.get('title', '').lower():
+                    results.append(document)
+        return results