Update app.py

app.py

@@ -8,10 +8,14 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter, SemanticChun
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.chains import ConversationalRetrievalChain
-from langchain.llms import ZhipuAI
-from langchain.prompts import PromptTemplate
-from langchain.memory import ConversationBufferMemory
 from langchain.schema import HumanMessage, AIMessage
+from langchain.memory import ConversationBufferMemory
+from langchain.prompts import PromptTemplate
+from langchain.llms.base import LLM
+from typing import Optional, List, Dict, Any
+
+# Import the new Zhipu AI SDK
+from zhipuai import ZhipuAI
 
 # Configuration from environment variables
 ZHIPU_MODEL = os.environ.get("ZHIPU_MODEL", "chatglm3-6b")
@@ -21,6 +25,37 @@ CHUNK_OVERLAP = int(os.environ.get("CHUNK_OVERLAP", "200"))
 USE_SEMANTIC_CHUNKING = os.environ.get("USE_SEMANTIC_CHUNKING", "true").lower() == "true"
 MEMORY_K = int(os.environ.get("MEMORY_K", "5"))  # Number of messages to remember
 
+# Custom LLM wrapper for the new Zhipu AI SDK
+class ZhipuAILLM(LLM):
+    client: ZhipuAI
+    model: str = "chatglm3-6b"
+    temperature: float = 0.1
+
+    def __init__(self, api_key: str, **kwargs):
+        super().__init__(**kwargs)
+        self.client = ZhipuAI(api_key=api_key)
+        self.model = kwargs.get("model", self.model)
+        self.temperature = kwargs.get("temperature", self.temperature)
+
+    @property
+    def _llm_type(self) -> str:
+        return "zhipuai"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> str:
+        response = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=self.temperature,
+            **kwargs
+        )
+        return response.choices[0].message.content
+
 # Database setup
 DB_PATH = "/data/chat_history.db"
 
@@ -118,10 +153,10 @@ def initialize_system(pdf_path):
     vector_store = FAISS.from_texts(texts, embeddings)
     retriever = vector_store.as_retriever()
     
-    # Initialize Zhipu LLM
-    llm = ZhipuAI(
-        model=ZHIPU_MODEL,
+    # Initialize Zhipu LLM using our custom wrapper
+    llm = ZhipuAILLM(
         api_key=os.environ["ZHIPU_API_KEY"],
+        model=ZHIPU_MODEL,
         temperature=ZHIPU_TEMPERATURE
     )