isage-middleware 0.1.4__py3-none-any.whl → 0.1.5.1__py3-none-any.whl

sage/middleware/utils/embedding/embedding_model.py

@@ -1,19 +1,32 @@
 import os
-import sys
-os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
-
-# 添加项目根目录到Python路径
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
 
+# flake8: noqa: E402
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+import sys
 import time
 
 from dotenv import load_dotenv
+from sage.middleware.utils.embedding import (
+    _cohere,
+    bedrock,
+    hf,
+    jina,
+    lollms,
+    mockembedder,
+    nvidia_openai,
+    ollama,
+    openai,
+    siliconcloud,
+    zhipu,
+)
+from transformers import AutoModel, AutoTokenizer
 
 load_dotenv()
 
-from sage.middleware.utils.embedding import hf, ollama, siliconcloud, openai, bedrock, zhipu, mockembedder  # , instructor
-from sage.middleware.utils.embedding import _cohere, nvidia_openai, lollms, jina
-from transformers import AutoModel, AutoTokenizer
+# Ensure project root is on sys.path for imports that rely on package layout
+sys.path.append(
+    os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
+)
 
 
 class EmbeddingModel:
@@ -29,13 +42,13 @@ class EmbeddingModel:
         if method == "default":
             method = "hf"
             kwargs["model"] = "sentence-transformers/all-MiniLM-L6-v2"
-            
+
         if method == "mockembedder":
             kwargs["model"] = "mockembedder"  # 确保 model 参数存在
             if "fixed_dim" not in kwargs:
                 kwargs["fixed_dim"] = 128  # 默认维度
-                
-        self.set_dim(kwargs["model"] )
+
+        self.set_dim(kwargs["model"])
         self.method = method
 
         # self.kwargs = {}
@@ -44,14 +57,24 @@ class EmbeddingModel:
             if "model" not in kwargs:
                 raise ValueError("hf method need model")
             model_name = kwargs["model"]
-            self.kwargs["tokenizer"] = AutoTokenizer.from_pretrained(model_name)
-            self.kwargs["embed_model"] = AutoModel.from_pretrained(model_name, trust_remote_code=True)
-            self.kwargs.pop("model")
+            # Load HF models - fail explicitly if unavailable
+            try:
+                self.kwargs["tokenizer"] = AutoTokenizer.from_pretrained(model_name)
+                self.kwargs["embed_model"] = AutoModel.from_pretrained(
+                    model_name, trust_remote_code=True
+                )
+                self.kwargs.pop("model")
+            except Exception as e:
+                # 明确失败，不静默回退到mockembedder
+                raise RuntimeError(
+                    f"Failed to load embedding model '{model_name}': {e}. "
+                    f"Please ensure the model is available or use a different embedding method. "
+                    f"For testing with mock embedder, explicitly set method='mockembedder'."
+                ) from e
         elif method == "mockembedder":
             # 初始化 mockembedder
             self.kwargs["embed_model"] = mockembedder.MockTextEmbedder(
-                model_name="mock-model",
-                fixed_dim=kwargs.get("fixed_dim", 128)
+                model_name="mock-model", fixed_dim=kwargs.get("fixed_dim", 128)
             )
         self.embed_fn = self._get_embed_function(method)
 
@@ -62,17 +85,17 @@ class EmbeddingModel:
         """
         dimension_mapping = {
             "mistral_embed": 1024,
-            "embed-multilingual-v3.0":1024,
+            "embed-multilingual-v3.0": 1024,
             "embed-english-v3.0": 1024,
             "embed-english-light-v3.0": 384,
             "embed-multilingual-light-v3.0": 384,
             "embed-english-v2.0": 4096,
             "embed-english-light-v2.0": 1024,
             "embed-multilingual-v2.0": 768,
-            "jina-embeddings-v3":1024,
-            "BAAI/bge-m3":1024,
-            "sentence-transformers/all-MiniLM-L6-v2":384,
-            "mockembedder": 128
+            "jina-embeddings-v3": 1024,
+            "BAAI/bge-m3": 1024,
+            "sentence-transformers/all-MiniLM-L6-v2": 384,
+            "mockembedder": 128,
         }
         if model_name in dimension_mapping:
             self.dim = dimension_mapping[model_name]
@@ -96,7 +119,9 @@ class EmbeddingModel:
             "ollama": ollama.ollama_embed_sync,
             "siliconcloud": siliconcloud.siliconcloud_embedding_sync,
             "cohere": _cohere.cohere_embed_sync,
-            "mockembedder": lambda text, **kwargs: kwargs["embed_model"].encode(text).tolist(),
+            "mockembedder": lambda text, **kwargs: kwargs["embed_model"]
+            .encode(text)
+            .tolist(),
             # "instructor": instructor.instructor_embed
         }
         if method not in mapping:
@@ -117,17 +142,17 @@ class EmbeddingModel:
 
     def embed(self, text: str) -> list[float]:
         return self._embed(text)
-    
+
     def encode(self, text: str) -> list[float]:
         return self._embed(text)
-    
+
     @property
     def method_name(self) -> str:
         """当前embedding方法名"""
         return self.init_method
 
 
-def apply_embedding_model(name: str = "default",**kwargs) -> EmbeddingModel:
+def apply_embedding_model(name: str = "default", **kwargs) -> EmbeddingModel:
     """
     usage  参见sage/api/model/operator_test.py
     while name(method) = "hf", please set the param:model;
@@ -135,16 +160,20 @@ def apply_embedding_model(name: str = "default",**kwargs) -> EmbeddingModel:
     while name(method) = "jina/siliconcloud/cohere",please set the params:api_key,model;
     Example:operator_test.py
     """
-    return EmbeddingModel(method=name,**kwargs)
+    return EmbeddingModel(method=name, **kwargs)
+
 
 def main():
-    embedding_model = EmbeddingModel(method="hf",model = "sentence-transformers/all-MiniLM-L6-v2")
+    embedding_model = EmbeddingModel(
+        method="hf", model="sentence-transformers/all-MiniLM-L6-v2"
+    )
     for i in range(10):
         start = time.time()
         v = embedding_model.embed(f"{i} times ")
         print(v)
         end = time.time()
         print(f"embedding time :{end-start}")
-        
-if __name__ =="__main__":
+
+
+if __name__ == "__main__":
     main()

sage/middleware/utils/embedding/hf.py

@@ -1,15 +1,15 @@
-import asyncio
-import copy
 import os
+
+# flake8: noqa: E402
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 from functools import lru_cache
-import asyncio
-import pipmaster as pm  # Pipmaster for dynamic library install
 
 # Dependencies should be installed via requirements.txt
 # transformers, torch, tenacity, and numpy are required for this module
 
 try:
-    from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
+    from transformers import AutoModel  # noqa: F401
+    from transformers import AutoModelForCausalLM, AutoTokenizer
 except ImportError:
     raise ImportError(
         "transformers package is required for HuggingFace embedding functionality. "
@@ -25,7 +25,7 @@ except ImportError:
     )
 
 try:
-    import tenacity
+    import tenacity  # noqa: F401
 except ImportError:
     raise ImportError(
         "tenacity package is required for HuggingFace embedding functionality. "
@@ -33,7 +33,7 @@ except ImportError:
     )
 
 try:
-    import numpy
+    import numpy  # noqa: F401
 except ImportError:
     raise ImportError(
         "numpy package is required for HuggingFace embedding functionality. "
@@ -42,6 +42,7 @@ except ImportError:
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+
 @lru_cache(maxsize=1)
 def initialize_hf_model(model_name):
     hf_tokenizer = AutoTokenizer.from_pretrained(
@@ -56,8 +57,6 @@ def initialize_hf_model(model_name):
     return hf_model, hf_tokenizer
 
 
-import torch
-
 def hf_embed_sync(text: str, tokenizer, embed_model) -> list[float]:
     """
     使用 HuggingFace 模型同步生成文本 embedding。
@@ -86,5 +85,3 @@ def hf_embed_sync(text: str, tokenizer, embed_model) -> list[float]:
         return embeddings.detach().to(torch.float32).cpu()[0].tolist()
     else:
         return embeddings.detach().cpu()[0].tolist()
-
-

sage/middleware/utils/embedding/instructor.py

@@ -1,10 +1,9 @@
 import numpy as np
-from InstructorEmbedding import INSTRUCTOR
 from sentence_transformers import SentenceTransformer
 
 
 async def instructor_embed(
-        texts: [str], model: str = "hkunlp/instructor-large"
+    texts: [str], model: str = "hkunlp/instructor-large"
 ) -> [np.array]:
     _model = SentenceTransformer(model)
     return _model.encode(texts)

sage/middleware/utils/embedding/jina.py

@@ -1,12 +1,12 @@
-import asyncio
 import os
-import pipmaster as pm  # Pipmaster for dynamic library install
+
+import requests
 
 # Dependencies should be installed via requirements.txt
 # tenacity is required for this module
 
 try:
-    import tenacity
+    import tenacity  # noqa: F401
 except ImportError:
     raise ImportError(
         "tenacity package is required for Jina embedding functionality. "
@@ -21,8 +21,6 @@ except ImportError:
         "Please install it via: pip install aiohttp"
     )
 
-import numpy as np
-
 
 async def fetch_data(url, headers, data):
     async with aiohttp.ClientSession() as session:
@@ -38,7 +36,7 @@ async def jina_embed(
     late_chunking: bool = False,
     base_url: str = None,
     api_key: str = None,
-    model:str="jina-embeddings-v3"
+    model: str = "jina-embeddings-v3",
 ) -> list[float]:
     if api_key:
         os.environ["JINA_API_KEY"] = api_key
@@ -57,10 +55,8 @@ async def jina_embed(
     }
     data_list = await fetch_data(url, headers, data)
     print(data_list)
-    return data_list[0]['embedding']
+    return data_list[0]["embedding"]
 
-import os
-import requests
 
 def jina_embed_sync(
     text: str,
@@ -68,7 +64,7 @@ def jina_embed_sync(
     late_chunking: bool = False,
     base_url: str = None,
     api_key: str = None,
-    model: str = "jina-embeddings-v3"
+    model: str = "jina-embeddings-v3",
 ) -> list[float]:
     """
     同步版本：调用 Jina AI embedding API 获取嵌入向量
@@ -108,8 +104,3 @@ def jina_embed_sync(
         return data["data"][0]["embedding"]
     except Exception as e:
         raise RuntimeError(f"Jina API call failed: {str(e)}")
-
-
-
-
-

sage/middleware/utils/embedding/lollms.py

@@ -1,10 +1,12 @@
 import sys
 
+import requests
+
 if sys.version_info < (3, 9):
-    from typing import AsyncIterator
+    pass
 else:
-    from collections.abc import AsyncIterator
-import pipmaster as pm  # Pipmaster for dynamic library install
+    pass
+
 
 # Dependencies should be installed via requirements.txt
 # aiohttp and tenacity are required for this module
@@ -18,7 +20,7 @@ except ImportError:
     )
 
 try:
-    import tenacity
+    import tenacity  # noqa: F401
 except ImportError:
     raise ImportError(
         "tenacity package is required for Lollms embedding functionality. "
@@ -26,7 +28,6 @@ except ImportError:
     )
 
 
-
 async def lollms_embed(
     text: str,
     embed_model=None,
@@ -62,7 +63,6 @@ async def lollms_embed(
             result = await response.json()
             return result["vector"]
 
-import requests
 
 def lollms_embed_sync(
     text: str,
@@ -92,7 +92,9 @@ def lollms_embed_sync(
     request_data = {"text": text}
 
     try:
-        response = requests.post(f"{base_url}/lollms_embed", json=request_data, headers=headers)
+        response = requests.post(
+            f"{base_url}/lollms_embed", json=request_data, headers=headers
+        )
         response.raise_for_status()
         result = response.json()
         return result["vector"]

sage/middleware/utils/embedding/mockembedder.py

@@ -1,25 +1,29 @@
-import torch
 import hashlib
 from typing import Optional
 
+import torch
+
+
 class MockTextEmbedder:
-    def __init__(self, model_name: str = 'mock-model', fixed_dim: int = 128):
+    def __init__(self, model_name: str = "mock-model", fixed_dim: int = 128):
         """Mock 模型，输出固定维度的随机张量（但相同文本输出一致）"""
         self.fixed_dim = fixed_dim
         # 用模型名作为随机种子的一部分，确保不同实例行为一致
         self.seed = int(hashlib.sha256(model_name.encode()).hexdigest()[:8], 16)
         # 添加 method_name 属性以兼容 MemoryManager
         self.method_name = "mockembedder"
-        
-    def encode(self, text: str, max_length: int = 512, stride: Optional[int] = None) -> torch.Tensor:
+
+    def encode(
+        self, text: str, max_length: int = 512, stride: Optional[int] = None
+    ) -> torch.Tensor:
         """生成固定维度的伪嵌入（相同文本输出相同）"""
         if not text.strip():
             return torch.zeros(self.fixed_dim)
-        
+
         # 根据文本内容生成确定性随机数
         text_seed = self.seed + int(hashlib.sha256(text.encode()).hexdigest()[:8], 16)
         torch.manual_seed(text_seed)
-        
+
         # 生成随机向量（与原有代码维度逻辑一致）
         if stride is None or len(text.split()) <= max_length:
             return self._embed_single()
@@ -39,8 +43,8 @@ class MockTextEmbedder:
     def _adjust_dimension(self, embedding: torch.Tensor) -> torch.Tensor:
         """保持与原代码一致的维度调整逻辑"""
         if embedding.size(0) > self.fixed_dim:
-            return embedding[:self.fixed_dim]
+            return embedding[: self.fixed_dim]
         elif embedding.size(0) < self.fixed_dim:
             padding = torch.zeros(self.fixed_dim - embedding.size(0))
             return torch.cat((embedding, padding))
-        return embedding
+        return embedding

sage/middleware/utils/embedding/nvidia_openai.py

@@ -1,25 +1,24 @@
-import asyncio
-import sys
 import os
+import sys
 
 if sys.version_info < (3, 9):
     pass
 else:
     pass
 
-import pipmaster as pm  # Pipmaster for dynamic library install
 
 # Dependencies should be installed via requirements.txt
 # openai is required for this module
 
 try:
-    from openai import AsyncOpenAI
+    from openai import AsyncOpenAI, OpenAI
 except ImportError:
     raise ImportError(
         "openai package is required for NVIDIA OpenAI embedding functionality. "
         "Please install it via: pip install openai"
     )
 
+
 async def nvidia_openai_embed(
     text: str,
     model: str = "nvidia/llama-3.2-nv-embedqa-1b-v1",
@@ -52,17 +51,14 @@ async def nvidia_openai_embed(
     return response.data[0].embedding
 
 
-import os
-from openai import OpenAI  # 同步 client
-
 def nvidia_openai_embed_sync(
     text: str,
     model: str = "nvidia/llama-3.2-nv-embedqa-1b-v1",
     base_url: str = "https://integrate.api.nvidia.com/v1",
     api_key: str = None,
     input_type: str = "passage",  # query for retrieval, passage for embedding
-    trunc: str = "NONE",          # NONE or START or END
-    encode: str = "float",        # float or base64
+    trunc: str = "NONE",  # NONE or START or END
+    encode: str = "float",  # float or base64
 ) -> list[float]:
     """
     同步版本：使用 NVIDIA NIM 接口生成文本 embedding。
@@ -88,10 +84,7 @@ def nvidia_openai_embed_sync(
         model=model,
         input=text,
         encoding_format=encode,
-        extra_body={
-            "input_type": input_type,
-            "truncate": trunc
-        },
+        extra_body={"input_type": input_type, "truncate": trunc},
     )
 
     return response.data[0].embedding

sage/middleware/utils/embedding/ollama.py

@@ -1,17 +1,16 @@
 import sys
 
 if sys.version_info < (3, 9):
-    from typing import AsyncIterator
+    pass
 else:
-    from collections.abc import AsyncIterator
+    pass
 
-import pipmaster as pm  # Pipmaster for dynamic library install
 
 # Dependencies should be installed via requirements.txt
 # ollama and tenacity are required for this module to work
 
 try:
-    import ollama
+    import ollama  # noqa: F401
 except ImportError:
     raise ImportError(
         "ollama package is required for Ollama embedding functionality. "
@@ -19,7 +18,7 @@ except ImportError:
     )
 
 try:
-    import tenacity
+    import tenacity  # noqa: F401
 except ImportError:
     raise ImportError(
         "tenacity package is required for Ollama embedding functionality. "
@@ -27,18 +26,7 @@ except ImportError:
     )
 
 
-
-import numpy as np
-from typing import Union
-
-
-
-
-async def ollama_embed(
-    text: str,
-    embed_model,
-    **kwargs
-) -> list:
+async def ollama_embed(text: str, embed_model, **kwargs) -> list:
     """
     Generate embedding for a single text using Ollama.
 
@@ -65,11 +53,8 @@ async def ollama_embed(
     data = ollama_client.embed(model=embed_model, input=text)
     return data["embedding"]
 
-def ollama_embed_sync(
-    text: str,
-    embed_model,
-    **kwargs
-) -> list[float]:
+
+def ollama_embed_sync(text: str, embed_model, **kwargs) -> list[float]:
     """
     同步版本：使用 Ollama 客户端生成 embedding 向量。
 

sage/middleware/utils/embedding/openai.py

@@ -1,27 +1,25 @@
-import asyncio
-import sys
 import os
-import logging
-
-
+import sys
 
 if sys.version_info < (3, 9):
-    from typing import AsyncIterator
+    pass
 else:
-    from collections.abc import AsyncIterator
-import pipmaster as pm  # Pipmaster for dynamic library install
+    pass
+
 
 # Dependencies should be installed via requirements.txt
 # openai is required for this module
 
 try:
-    from openai import AsyncOpenAI  # 确保导入了这个
+    from openai import AsyncOpenAI, OpenAI  # 确保导入了这个
+    from openai.types import CreateEmbeddingResponse
 except ImportError:
     raise ImportError(
         "openai package is required for OpenAI embedding functionality. "
         "Please install it via: pip install openai"
     )
 
+
 async def openai_embed(
     text: str,
     model: str = "text-embedding-3-small",
@@ -57,18 +55,12 @@ async def openai_embed(
     )
 
     response = await openai_async_client.embeddings.create(
-        model=model,
-        input=text,
-        encoding_format="float"
+        model=model, input=text, encoding_format="float"
     )
 
     return response.data[0].embedding
 
 
-import os
-from openai import OpenAI  # 同步版 client
-from openai.types import CreateEmbeddingResponse
-
 def openai_embed_sync(
     text: str,
     model: str = "text-embedding-3-small",
@@ -98,15 +90,11 @@ def openai_embed_sync(
     openai_sync_client = (
         OpenAI(default_headers=default_headers, api_key=api_key)
         if base_url is None
-        else OpenAI(
-            base_url=base_url, default_headers=default_headers, api_key=api_key
-        )
+        else OpenAI(base_url=base_url, default_headers=default_headers, api_key=api_key)
     )
 
     response: CreateEmbeddingResponse = openai_sync_client.embeddings.create(
-        model=model,
-        input=text,
-        encoding_format="float"
+        model=model, input=text, encoding_format="float"
     )
 
     return response.data[0].embedding

sage/middleware/utils/embedding/siliconcloud.py

@@ -1,40 +1,25 @@
-import asyncio
-import os
+import base64
+import struct
 import sys
 
+import requests
+
 if sys.version_info < (3, 9):
     pass
 else:
     pass
-import pipmaster as pm  # Pipmaster for dynamic library install
-
-
-from openai import (
-    APIConnectionError,
-    RateLimitError,
-    APITimeoutError,
-)
-from tenacity import (
-    retry,
-    stop_after_attempt,
-    wait_exponential,
-    retry_if_exception_type,
-)
-
-import numpy as np
-import aiohttp
 import base64
 import struct
 
-
+import aiohttp
 
 
 async def siliconcloud_embedding(
-        text: str,
-        model: str = "netease-youdao/bce-embedding-base_v1",
-        base_url: str = "https://api.siliconflow.cn/v1/embeddings",
-        max_token_size: int = 512,
-        api_key: str = None,
+    text: str,
+    model: str = "netease-youdao/bce-embedding-base_v1",
+    base_url: str = "https://api.siliconflow.cn/v1/embeddings",
+    max_token_size: int = 512,
+    api_key: str = None,
 ) -> list:
     """
     Generate embedding for a single text using SiliconCloud (NetEase Youdao).
@@ -76,9 +61,6 @@ async def siliconcloud_embedding(
     float_array = struct.unpack("<" + "f" * n, decode_bytes)
     return list(float_array)
 
-import requests
-import base64
-import struct
 
 def siliconcloud_embedding_sync(
     text: str,