Sync from GitHub Actions

Lawverse/agents/tools.py

@@ -7,13 +7,13 @@ from langchain_core.documents import Document
 def retrieve_with_hybrid_tool(retriever, query: str, top_k: int = 5) -> List[Document]:
     if retriever is None:
         return []
-    try:
+
+    if hasattr(retriever, "invoke"):
         docs = retriever.invoke(query)
-    except Exception:
-        try:
-            docs = retriever.get_relevant_documents(query)
-        except Exception:
-            docs = retriever._get_relevant_documents(query)
+    elif hasattr(retriever, "get_relevant_documents"):
+        docs = retriever.get_relevant_documents(query)
+    else:
+        docs = retriever._get_relevant_documents(query)
 
     return list(docs or [])[:top_k]
 

Lawverse/memory/langchain_memory.py

@@ -1,28 +1,42 @@
 import sys
 from flask import session, has_request_context
 from datetime import datetime
-from langchain_classic.memory import ConversationBufferMemory
+from langchain_core.messages import HumanMessage, AIMessage
 from Lawverse.logger import logging
 from Lawverse.exception import ExceptionHandle
 from Lawverse.storage.factory import get_chat_store
 
 
+class _SimpleChatHistory:
+    def __init__(self):
+        self.messages = []
+
+    def add_user_message(self, content: str):
+        self.messages.append(HumanMessage(content=content or ""))
+
+    def add_ai_message(self, content: str):
+        self.messages.append(AIMessage(content=content or ""))
+
+    def clear(self):
+        self.messages.clear()
+
+
+class _SimpleMemory:
+    def __init__(self):
+        self.chat_memory = _SimpleChatHistory()
+
+    def clear(self):
+        self.chat_memory.clear()
+
 class ChatMemory:
     def __init__(self, chat_id=None, user_id=None):
         try:
             self.chat_id = chat_id or self._create_new_chat_id()
             self.user_id = str(user_id or self._get_current_user_id())
             self.store = get_chat_store()
-
-            self.memory = ConversationBufferMemory(
-                memory_key="chat_history",
-                return_messages=True,
-                output_key="answer",
-            )
-
+            self.memory = _SimpleMemory()
             self._load_memory()
             logging.info(f"ChatMemory initialized for user_id={self.user_id}, chat_id={self.chat_id}")
-
         except Exception as e:
             raise ExceptionHandle(e, sys) from e
 
@@ -37,22 +51,19 @@ class ChatMemory:
     def _load_memory(self):
         try:
             data = self.store.load_chat(self.user_id, self.chat_id)
-
             if not data:
                 logging.info(f"No existing memory found for chat_id: {self.chat_id}")
                 return
 
-            for msg in data.get("history", []):
+            for msg in data.get("history", []) or []:
                 user_msg = msg.get("user", "")
                 ai_msg = msg.get("ai", "")
-
                 if user_msg:
                     self.memory.chat_memory.add_user_message(user_msg)
                 if ai_msg:
                     self.memory.chat_memory.add_ai_message(ai_msg)
 
             logging.info(f"Memory loaded successfully for chat_id: {self.chat_id}")
-
         except Exception as e:
             raise ExceptionHandle(e, sys) from e
 
@@ -65,19 +76,16 @@ class ChatMemory:
     def _history_as_pairs(self):
         messages = self.memory.chat_memory.messages
         history = []
-
         i = 0
         while i < len(messages):
             user_msg = messages[i].content if i < len(messages) else ""
             ai_msg = messages[i + 1].content if i + 1 < len(messages) else ""
-
             if user_msg or ai_msg:
                 history.append({"user": user_msg, "ai": ai_msg})
             i += 2
-
         return history
 
-    def save_memory(self):
+    def save_memory(self) -> bool:
         try:
             self.store.save_chat(
                 user_id=self.user_id,
@@ -85,14 +93,17 @@ class ChatMemory:
                 title=self._get_title(),
                 history=self._history_as_pairs(),
             )
-
             logging.info(f"Memory saved successfully for chat_id: {self.chat_id}")
+            return True
         except Exception as e:
-            raise ExceptionHandle(e, sys) from e
+            logging.error(f"Memory save failed for chat_id={self.chat_id}: {e}")
+            return False
 
     def _get_title(self):
-        if self.memory.chat_memory.messages:
-            return self.memory.chat_memory.messages[0].content[:40]
+        for msg in self.memory.chat_memory.messages:
+            content = getattr(msg, "content", "") or ""
+            if content.strip():
+                return content.strip()[:40]
         return "New legal query"
 
     def clear_memory(self):
@@ -100,6 +111,5 @@ class ChatMemory:
             self.memory.clear()
             self.store.delete_chat(self.user_id, self.chat_id)
             logging.info(f"Memory deleted for chat_id: {self.chat_id}")
-
         except Exception as e:
             raise ExceptionHandle(e, sys) from e

Lawverse/retrieval/hybrid.py

@@ -37,7 +37,7 @@ def hybrid_retrieve(
 ):
     try:
         dense_results = faiss_db.similarity_search(query, k=initial_top_k)
-        sparse_results = bm25_retrieve(bm25, chunks, query, top_k=initial_top_k)
+        sparse_results = bm25_retrieve(bm25, query, chunks, top_k=initial_top_k)
 
         doc_map: Dict[str, Document] = {}
         rrf_scores = defaultdict(float)

Lawverse/retrieval/sparse.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import re
 import sys
-from typing import List, Tuple, Union
+from typing import List, Tuple, Union, Iterable
 from rank_bm25 import BM25Okapi
 from langchain_core.documents import Document
 from Lawverse.logger import logging
@@ -18,41 +18,34 @@ STOP_WORDS = {
 }
 
 
-def bm25_tokenizer(text: Union[str, List[str], tuple]) -> List[str]:
+def bm25_tokenizer(text: Union[str, Iterable[str], None]) -> List[str]:
     if text is None:
         return []
 
-    if isinstance(text, (list, tuple)):
-        return [
-            str(token).lower().strip()
-            for token in text
-            if str(token).strip() and str(token).lower().strip() not in STOP_WORDS
-        ]
+    if isinstance(text, (list, tuple, set)):
+        tokens = [str(token).lower().strip() for token in text]
+        return [token for token in tokens if len(token) > 1 and token not in STOP_WORDS]
 
     text = str(text).lower()
     english_tokens = re.findall(r"[a-zA-Z0-9]+", text)
     bangla_tokens = re.findall(r"[\u0980-\u09FF]+", text)
     tokens = english_tokens + bangla_tokens
-
-    return [
-        token
-        for token in tokens
-        if len(token) > 1 and token not in STOP_WORDS
-    ]
+    return [token for token in tokens if len(token) > 1 and token not in STOP_WORDS]
 
 
 def build_sparse_index(chunks: List[Document], k1: float = 1.5, b: float = 0.8) -> BM25Okapi:
     try:
         logging.info("Building sparse BM25 index...")
-        tokenized_corpus = [
-            bm25_tokenizer(chunk.page_content)
-            for chunk in chunks
-        ]
-
+        safe_chunks = [chunk for chunk in chunks if isinstance(chunk, Document)]
+        if len(safe_chunks) != len(chunks):
+            logging.warning(
+                "BM25 received non-Document chunks; ignored %s invalid items.",
+                len(chunks) - len(safe_chunks),
+            )
+        tokenized_corpus = [bm25_tokenizer(chunk.page_content) for chunk in safe_chunks]
         bm25 = BM25Okapi(tokenized_corpus, k1=k1, b=b)
-        logging.info(f"BM25 sparse index successfully built with {len(chunks)} chunks.")
+        logging.info(f"BM25 sparse index successfully built with {len(safe_chunks)} chunks.")
         return bm25
-
     except Exception as e:
         logging.error(f"Failed to build BM25 sparse index. Error: {e}")
         raise ExceptionHandle(e, sys)
@@ -65,24 +58,28 @@ def bm25_retrieve(
     top_k: int = 10,
 ) -> List[Tuple[Document, float]]:
     try:
+        if not isinstance(chunks, list):
+            raise TypeError(
+                f"chunks must be a list[Document], got {type(chunks).__name__}. "
+            )
+
         query_tokens = bm25_tokenizer(query)
         scores = bm25.get_scores(query_tokens)
+        ranked = sorted(enumerate(scores), key=lambda item: item[1], reverse=True)[:top_k]
 
-        ranked = sorted(
-            enumerate(scores),
-            key=lambda item: item[1],
-            reverse=True,
-        )[:top_k]
-
-        results = []
+        results: List[Tuple[Document, float]] = []
         for idx, score in ranked:
+            if idx >= len(chunks):
+                continue
             doc = chunks[idx]
-            doc.metadata = dict(doc.metadata or {})
+            if not isinstance(doc, Document):
+                logging.warning("Skipping BM25 result with invalid chunk type: %s", type(doc).__name__)
+                continue
+            doc = Document(page_content=doc.page_content, metadata=dict(doc.metadata or {}))
             doc.metadata["bm25_score"] = float(score)
             results.append((doc, float(score)))
 
         return results
-
     except Exception as e:
         logging.error(f"BM25 retrieval failed. Error: {e}")
         raise ExceptionHandle(e, sys)

api/app.py

@@ -1,8 +1,8 @@
 from flask import Flask, render_template, request, jsonify, session, stream_with_context, Response
 from dotenv import load_dotenv
-import secrets
-import os
 from threading import Lock
+import os
+import secrets
 import logging as py_logging
 from Lawverse.pipeline.rag_pipeline import rag_components
 from Lawverse.pipeline.llm_loader import llm
@@ -13,15 +13,12 @@ from Lawverse.agents.graph import create_agentic_chain
 from Lawverse.storage.factory import get_chat_store
 from api.auth import auth_bp, login_required
 for logger_name in [
-    "httpcore",
-    "httpx",
-    "hpack",
-    "filelock",
-    "sentence_transformers",
-    "urllib3",
+    "httpcore", "httpx", "hpack", "filelock", "sentence_transformers",
+    "urllib3", "openai", "openai._base_client", "faiss", "datasets",
 ]:
     py_logging.getLogger(logger_name).setLevel(py_logging.WARNING)
-    
+
+
 load_dotenv()
 app = Flask(__name__, template_folder="../templates")
 app.secret_key = os.getenv("SECRET_KEY") or secrets.token_hex(32)
@@ -31,11 +28,12 @@ app.register_blueprint(monitor_bp)
 
 BASE_COMPONENTS = None
 BASE_COMPONENTS_LOCK = Lock()
+AGENT_CHAIN = None
+AGENT_CHAIN_LOCK = Lock()
 active_chains = {}
 
 def get_base_components():
     global BASE_COMPONENTS
-    
     if BASE_COMPONENTS is not None:
         return BASE_COMPONENTS
 
@@ -44,19 +42,28 @@ def get_base_components():
             logging.info("Loading Lawverse RAG base components...")
             BASE_COMPONENTS = rag_components()
             logging.info("Lawverse RAG base components loaded successfully.")
-
     return BASE_COMPONENTS
 
 
-def create_agent_session(chat_id=None):
-    components = get_base_components()
-    chain = create_agentic_chain(components, llm)
-    memory_manager = ChatMemory(chat_id=chat_id)
+def get_agent_chain():
+    global AGENT_CHAIN
+    if AGENT_CHAIN is not None:
+        return AGENT_CHAIN
 
+    with AGENT_CHAIN_LOCK:
+        if AGENT_CHAIN is None:
+            components = get_base_components()
+            AGENT_CHAIN = create_agentic_chain(components, llm)
+    return AGENT_CHAIN
+
+
+def create_agent_session(chat_id=None, save_on_create=True):
+    chain = get_agent_chain()
+    memory_manager = ChatMemory(chat_id=chat_id)
     active_chains[memory_manager.chat_id] = (chain, memory_manager)
     session["chat_id"] = memory_manager.chat_id
-    memory_manager.save_memory()
-
+    if save_on_create:
+        memory_manager.save_memory()
     return chain, memory_manager
 
 @app.route("/", methods=["GET"])
@@ -69,18 +76,13 @@ def chat():
     chat_id = session.get("chat_id")
     if not chat_id or chat_id not in active_chains:
         create_agent_session()
-        
     return render_template("chat.html")
 
 @app.route("/new_chat", methods=["POST"])
 @login_required
 def new_chat():
     _, memory_manager = create_agent_session()
-
-    return jsonify({
-        "chat_id": memory_manager.chat_id,
-        "title": memory_manager._get_title()
-    })
+    return jsonify({"chat_id": memory_manager.chat_id, "title": memory_manager._get_title()})
 
 
 @app.route("/response", methods=["POST"])
@@ -95,12 +97,12 @@ def rag_response():
         qa, memory_manager = active_chains[chat_id]
         data = request.get_json(silent=True) or {}
         query = data.get("message", "").strip()
-
         if not query:
             return jsonify({"error": "Empty message"}), 400
 
         def generate():
             answer_parts = []
+            generation_failed = False
             try:
                 for chunk in qa.stream({
                     "input": query,
@@ -109,15 +111,15 @@ def rag_response():
                     text = chunk if isinstance(chunk, str) else str(chunk)
                     answer_parts.append(text)
                     yield text
+            except Exception as e:
+                generation_failed = True
+                logging.error(f"Error during answer generation: {e}")
 
+            if not generation_failed:
                 full_answer = "".join(answer_parts).strip()
                 memory_manager.append_exchange(query, full_answer)
                 memory_manager.save_memory()
 
-            except Exception as e:
-                logging.error(f"Error during stream generation: {e}")
-                yield "**Error:** An error occurred while processing your request."
-
         return Response(stream_with_context(generate()), mimetype="text/plain")
 
     except Exception as e:
@@ -132,7 +134,6 @@ def get_chats():
         user_id = str(session.get("user_id"))
         chats = get_chat_store().list_chats(user_id)
         return jsonify(chats), 200
-
     except Exception as e:
         logging.error(f"Failed to list chats from cloud store: {e}")
         return jsonify([]), 200
@@ -147,27 +148,20 @@ def load_chat(chat_id):
         if not data:
             return jsonify({"error": "Chat not found"}), 404
 
-        _, memory_manager = create_agent_session(chat_id=chat_id)
-
+        _, memory_manager = create_agent_session(chat_id=chat_id, save_on_create=False)
         messages_list = memory_manager.memory.chat_memory.messages
         messages = []
-
         for i in range(0, len(messages_list), 2):
             user_msg = messages_list[i].content if i < len(messages_list) else None
             ai_msg = messages_list[i + 1].content if i + 1 < len(messages_list) else ""
-
             if user_msg:
-                messages.append({
-                    "user": user_msg,
-                    "ai": ai_msg
-                })
+                messages.append({"user": user_msg, "ai": ai_msg})
 
         return jsonify({
             "chat_id": chat_id,
             "title": memory_manager._get_title(),
             "messages": messages,
         }), 200
-
     except Exception as e:
         logging.error(f"Failed to load chat from cloud store: {e}")
         return jsonify({"error": "Internal Server Error"}), 500
@@ -178,21 +172,13 @@ def load_chat(chat_id):
 def delete_chat(chat_id):
     try:
         user_id = str(session.get("user_id"))
-
         deleted = get_chat_store().delete_chat(user_id, chat_id)
-
         was_active = chat_id in active_chains
         if was_active:
             del active_chains[chat_id]
-
         if session.get("chat_id") == chat_id:
             session.pop("chat_id", None)
-
-        return jsonify({
-            "success": deleted,
-            "was_active": was_active
-        }), 200
-
+        return jsonify({"success": deleted, "was_active": was_active}), 200
     except Exception as e:
         logging.error(f"Error deleting cloud chat {chat_id}: {e}")
         return jsonify({"error": "Internal Server Error"}), 500