Merge pull request #31 from kkkjz/main

fix model name

Author: kkkjz

memoryos-chromadb/comprehensive_test.py

@@ -21,14 +21,15 @@ def main():
     memoryos = Memoryos(
         user_id='travel_user_test',
         openai_api_key='',
-        openai_base_url='',
+        openai_base_url='https://cn2us02.opapi.win/v1',
         data_storage_path='./comprehensive_test_data',
         assistant_id='travel_assistant',
-        embedding_model_name='BAAI/bge-m3',
+        embedding_model_name='',
         mid_term_capacity=1000,
-        mid_term_heat_threshold=13.0,
+        mid_term_heat_threshold=12.0,
         mid_term_similarity_threshold=0.7,
-        short_term_capacity=2
+        short_term_capacity=2,
+        llm_model='gpt-4.1-mini'
     )
 
     print("📝 Phase 1: Adding 30 rounds of travel planning conversations...")

memoryos-chromadb/long_term.py

@@ -15,12 +15,14 @@ def __init__(self,
                  llm_interface: OpenAIClient,
                  knowledge_capacity=100, 
                  embedding_model_name: str = "all-MiniLM-L6-v2", 
-                 embedding_model_kwargs: Optional[dict] = None):
+                 embedding_model_kwargs: Optional[dict] = None,
+                 llm_model: str = "gpt-4o-mini"):  # 添加 llm_model 参数
         self.storage = storage_provider
         self.llm_interface = llm_interface
         self.knowledge_capacity = knowledge_capacity
         self.embedding_model_name = embedding_model_name
         self.embedding_model_kwargs = embedding_model_kwargs or {}
+        self.llm_model = llm_model  # 保存模型名称
 
     def update_user_profile(self, user_id: str, conversation_history: str) -> Optional[Dict[str, Any]]:
         """
@@ -31,6 +33,7 @@ def update_user_profile(self, user_id: str, conversation_history: str) -> Option
         updated_profile = gpt_user_profile_analysis(
             conversation_str=conversation_history,
             client=self.llm_interface,
+            model=self.llm_model,  # 传递模型参数
             existing_user_profile=existing_profile_str
         )
 
@@ -73,7 +76,11 @@ def extract_knowledge_from_text(self, text: str) -> Optional[Dict[str, Any]]:
         """
         if not text.strip():
             return None
-        return gpt_knowledge_extraction(conversation_str=text, client=self.llm_interface)
+        return gpt_knowledge_extraction(
+            conversation_str=text, 
+            client=self.llm_interface,
+            model=self.llm_model  # 传递模型参数
+        )
 
     def get_user_knowledge(self) -> list:
         return self.storage.get_all_user_knowledge()
@@ -95,4 +102,4 @@ def search_knowledge(self, query: str, knowledge_type: str = "user", top_k=5) ->
             results = self.storage.search_assistant_knowledge(query_vec, top_k=top_k)
 
         print(f"LongTermMemory: Searched {knowledge_type} knowledge for '{query[:30]}...'. Found {len(results)} matches.")
-        return results 
+        return results

memoryos-chromadb/memoryos.py

@@ -91,21 +91,24 @@ def __init__(self, user_id: str,
             client=self.client, 
             max_capacity=mid_term_capacity,
             embedding_model_name=self.embedding_model_name,
-            embedding_model_kwargs=self.embedding_model_kwargs
+            embedding_model_kwargs=self.embedding_model_kwargs,
+            llm_model=self.llm_model
         )
         self.user_long_term_memory = LongTermMemory(
             storage_provider=self.storage_provider,
             llm_interface=self.client,
             embedding_model_name=self.embedding_model_name,
-            embedding_model_kwargs=self.embedding_model_kwargs
+            embedding_model_kwargs=self.embedding_model_kwargs,
+            llm_model=self.llm_model 
         )
 
         # Initialize Memory Module for Assistant Knowledge
         self.assistant_long_term_memory = LongTermMemory(
             storage_provider=self.storage_provider,
             llm_interface=self.client,
             embedding_model_name=self.embedding_model_name,
-            embedding_model_kwargs=self.embedding_model_kwargs
+            embedding_model_kwargs=self.embedding_model_kwargs,
+            llm_model=self.llm_model
         )
 
         # Initialize Orchestration Modules

memoryos-chromadb/mid_term.py

@@ -43,11 +43,13 @@ def __init__(self,
                  client: OpenAIClient, 
                  max_capacity=2000,
                  embedding_model_name: str = "all-MiniLM-L6-v2", 
-                 embedding_model_kwargs: Optional[dict] = None):
+                 embedding_model_kwargs: Optional[dict] = None,
+                 llm_model: str = "gpt-4o-mini"):
         self.user_id = user_id
         self.client = client
         self.max_capacity = max_capacity
         self.storage = storage_provider
+        self.llm_model = llm_model
 
         # Load sessions and other data from the shared storage provider's in-memory metadata
         self.sessions: dict = self.storage.get_mid_term_sessions()
@@ -100,7 +102,7 @@ def add_session(self, summary, details):
             **self.embedding_model_kwargs
         )
         summary_vec = normalize_vector(summary_vec).tolist()
-        summary_keywords = list(extract_keywords_from_multi_summary(summary, client=self.client))
+        summary_keywords = list(extract_keywords_from_multi_summary(summary, client=self.client,model=self.llm_model))  
 
         processed_details = []
         for page_data in details:
@@ -132,7 +134,7 @@ def add_session(self, summary, details):
             else:
                 print(f"MidTermMemory: Computing new keywords for page {page_id}")
                 full_text = f"User: {page_data.get('user_input','')} Assistant: {page_data.get('agent_response','')}"
-                page_keywords = list(extract_keywords_from_multi_summary(full_text, client=self.client))
+                page_keywords = list(extract_keywords_from_multi_summary(full_text, client=self.client,model=self.llm_model))
 
             processed_page = {
                 **page_data, # Carry over existing fields like user_input, agent_response, timestamp
@@ -249,7 +251,7 @@ def insert_pages_into_session(self, summary_for_new_pages, keywords_for_new_page
 
                 if "page_keywords" not in page_data or not page_data["page_keywords"]:
                     full_text = f"User: {page_data.get('user_input','')} Assistant: {page_data.get('agent_response','')}"
-                    page_data["page_keywords"] = list(extract_keywords_from_multi_summary(full_text, client=self.client))
+                    page_data["page_keywords"] = list(extract_keywords_from_multi_summary(full_text, client=self.client,model=self.llm_model))
 
                 processed_new_pages.append({**page_data, "page_id": page_id})
 
@@ -285,7 +287,7 @@ def search_sessions(self, query_text, segment_similarity_threshold=0.1, page_sim
             **self.embedding_model_kwargs
         )
         query_vec = normalize_vector(query_vec)
-        query_keywords = set(extract_keywords_from_multi_summary(query_text, client=self.client))
+        query_keywords = set(extract_keywords_from_multi_summary(query_text, client=self.client,model=self.llm_model))
 
         # Search sessions using ChromaDB
         similar_sessions = self.storage.search_mid_term_sessions(query_vec.tolist(), top_k=top_k_sessions)

memoryos-chromadb/updater.py

@@ -53,7 +53,7 @@ def _process_page_embedding_and_keywords(self, page_data):
             tasks.append(('embedding', lambda: get_embedding(full_text)))
 
         if not ("page_keywords" in page_data and page_data["page_keywords"]):
-            tasks.append(('keywords', lambda: extract_keywords_from_multi_summary(full_text, client=self.client)))
+            tasks.append(('keywords', lambda: extract_keywords_from_multi_summary(full_text, client=self.client,model=self.llm_model)))
 
         if tasks:
             with ThreadPoolExecutor(max_workers=2) as executor: