mem1 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

mem1/__init__.py

@@ -1,16 +1,17 @@
 """
-Mem1 - 基于 Elasticsearch 的用户记忆系统
+Mem1 - 用户记忆系统（支持可插拔存储后端）
 """
 import logging
 
-__version__ = "0.0.5"
+__version__ = "0.0.7"
 
 # 屏蔽第三方库的详细日志（必须在导入前设置）
 logging.getLogger("elastic_transport").setLevel(logging.WARNING)
 logging.getLogger("elastic_transport.transport").setLevel(logging.WARNING)
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
-from mem1.memory_es import Mem1Memory
+from mem1.memory import Mem1Memory
 from mem1.config import Mem1Config, LLMConfig
+from mem1.storage import StorageBackend, ESStorage
 
-__all__ = ["Mem1Memory", "Mem1Config", "LLMConfig"]
+__all__ = ["Mem1Memory", "Mem1Config", "LLMConfig", "StorageBackend", "ESStorage"]

mem1/config.py

@@ -13,14 +13,21 @@ class LLMConfig(BaseModel):
 
 
 class VLConfig(BaseModel):
-    """视觉语言模型配置（可选，配置了 model 即启用，使用 dashscope SDK）"""
+    """视觉语言模型配置（可选，配置了 provider 即启用）
+    
+    支持的 provider:
+    - qwen: 使用 dashscope SDK 调用 Qwen-VL
+    - doubao: 使用 OpenAI 兼容接口调用豆包视觉模型
+    """
+    provider: str = ""  # qwen / doubao
     model: str = ""
     api_key: str = ""
+    base_url: str = ""  # doubao 需要
     
     @property
     def enabled(self) -> bool:
-        """只要配置了 model 就启用"""
-        return bool(self.model)
+        """配置了 provider 就启用"""
+        return bool(self.provider)
 
 
 class MemoryConfig(BaseModel):
@@ -101,10 +108,12 @@ class Mem1Config(BaseModel):
         memory_dir = required_vars["MEM1_MEMORY_DIR"]
         images_dir = f"{memory_dir}/images"
         
-        # VL 模型配置（可选，配置了 model 即启用，使用 dashscope SDK）
+        # VL 模型配置（可选，配置了 provider 即启用）
         vl_config = VLConfig(
+            provider=os.getenv("MEM1_VL_PROVIDER", ""),
             model=os.getenv("MEM1_VL_MODEL", ""),
-            api_key=os.getenv("MEM1_VL_API_KEY", "")
+            api_key=os.getenv("MEM1_VL_API_KEY", ""),
+            base_url=os.getenv("MEM1_VL_BASE_URL", "")
         )
         
         return cls(

mem1/llm.py

@@ -44,12 +44,20 @@ class LLMClient:
 
 
 class VLClient:
-    """视觉语言模型客户端（基于 dashscope SDK）"""
+    """视觉语言模型客户端（支持 qwen/doubao）"""
     
     def __init__(self, config: VLConfig):
         self.config = config
-        import dashscope
-        dashscope.api_key = config.api_key
+        self.provider = config.provider.lower()
+        
+        if self.provider == "qwen":
+            import dashscope
+            dashscope.api_key = config.api_key
+        elif self.provider == "doubao":
+            self.client = OpenAI(
+                api_key=config.api_key,
+                base_url=config.base_url
+            )
     
     def understand_image(
         self,
@@ -65,13 +73,22 @@ class VLClient:
         Returns:
             图片理解结果（包含 OCR 文字和内容描述）
         """
-        import dashscope
-        
         prompt = "请分析这张图片，完成以下任务：\n1. OCR识别：提取图片中的所有文字\n2. 内容理解：描述图片的主要内容和关键信息\n\n请用简洁的中文回答，格式如下：\n【文字内容】...\n【图片描述】..."
         
         if user_description:
             prompt += f"\n\n用户补充说明：{user_description}"
         
+        if self.provider == "qwen":
+            return self._call_qwen(image_path, prompt)
+        elif self.provider == "doubao":
+            return self._call_doubao(image_path, prompt)
+        else:
+            raise ValueError(f"不支持的 VL provider: {self.provider}")
+    
+    def _call_qwen(self, image_path: str, prompt: str) -> str:
+        """调用 Qwen-VL（dashscope SDK）"""
+        import dashscope
+        
         messages = [
             {
                 "role": "user",
@@ -88,3 +105,35 @@ class VLClient:
         )
         
         return response.output.choices[0].message.content[0]["text"]
+    
+    def _call_doubao(self, image_path: str, prompt: str) -> str:
+        """调用豆包视觉模型（OpenAI 兼容接口）"""
+        import base64
+        
+        # 读取图片并转为 base64
+        with open(image_path, "rb") as f:
+            image_data = base64.b64encode(f.read()).decode("utf-8")
+        
+        # 获取图片格式
+        ext = image_path.lower().split(".")[-1]
+        mime_type = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg", "gif": "image/gif", "webp": "image/webp"}.get(ext, "image/png")
+        
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{mime_type};base64,{image_data}"}
+                    },
+                    {"type": "text", "text": prompt}
+                ]
+            }
+        ]
+        
+        response = self.client.chat.completions.create(
+            model=self.config.model,
+            messages=messages
+        )
+        
+        return response.choices[0].message.content