isage-middleware 0.1.0__py3-none-any.whl

sage/middleware/utils/embedding/bedrock.py

@@ -0,0 +1,174 @@
+import copy
+import os
+import json
+
+import pipmaster as pm  # Pipmaster for dynamic library install
+
+# Dependencies should be installed via requirements.txt
+# aioboto3 and tenacity are required for this module
+
+try:
+    import aioboto3
+except ImportError:
+    raise ImportError(
+        "aioboto3 package is required for AWS Bedrock embedding functionality. "
+        "Please install it via: pip install aioboto3"
+    )
+
+try:
+    from tenacity import (
+        retry,
+        stop_after_attempt,
+        wait_exponential,
+        retry_if_exception_type,
+    )
+except ImportError:
+    raise ImportError(
+        "tenacity package is required for AWS Bedrock embedding functionality. "
+        "Please install it via: pip install tenacity"
+    )
+
+import numpy as np
+
+
+
+class BedrockError(Exception):
+    """Generic error for issues related to Amazon Bedrock"""
+
+
+
+async def bedrock_embed(
+    text: str,
+    model: str = "amazon.titan-embed-text-v2:0",
+    aws_access_key_id=None,
+    aws_secret_access_key=None,
+    aws_session_token=None,
+) -> list:
+    os.environ["AWS_ACCESS_KEY_ID"] = os.environ.get(
+        "AWS_ACCESS_KEY_ID", aws_access_key_id
+    )
+    os.environ["AWS_SECRET_ACCESS_KEY"] = os.environ.get(
+        "AWS_SECRET_ACCESS_KEY", aws_secret_access_key
+    )
+    os.environ["AWS_SESSION_TOKEN"] = os.environ.get(
+        "AWS_SESSION_TOKEN", aws_session_token
+    )
+
+    session = aioboto3.Session()
+    async with session.client("bedrock-sage.kernels.runtime") as bedrock_async_client:
+        model_provider = model.split(".")[0]
+
+        if model_provider == "amazon":
+            if "v2" in model:
+                body = json.dumps({
+                    "inputText": text,
+                    "embeddingTypes": ["float"],
+                })
+            elif "v1" in model:
+                body = json.dumps({"inputText": text})
+            else:
+                raise ValueError(f"Model {model} is not supported!")
+
+            response = await bedrock_async_client.invoke_model(
+                modelId=model,
+                body=body,
+                accept="application/json",
+                contentType="application/json",
+            )
+
+            response_body = await response.get("body").json()
+            return response_body["embedding"]
+
+        elif model_provider == "cohere":
+            body = json.dumps({
+                "texts": [text],
+                "input_type": "search_document",
+                "truncate": "NONE",
+            })
+
+            response = await bedrock_async_client.invoke_model(
+                model=model,
+                body=body,
+                accept="application/json",
+                contentType="application/json",
+            )
+
+            response_body = json.loads(response.get("body").read())
+            return response_body["embeddings"][0]
+
+        else:
+            raise ValueError(f"Model provider '{model_provider}' is not supported!")
+
+
+import os
+import json
+import boto3
+
+def bedrock_embed_sync(
+    text: str,
+    model: str = "amazon.titan-embed-text-v2:0",
+    aws_access_key_id=None,
+    aws_secret_access_key=None,
+    aws_session_token=None,
+) -> list[float]:
+    """
+    同步版本：使用 AWS Bedrock 生成 embedding。
+
+    Args:
+        text: 输入文本
+        model: 模型 ID，例如 "amazon.titan-embed-text-v2:0"
+        aws_access_key_id / secret / session_token: 可选 AWS 认证信息
+
+    Returns:
+        list[float]: embedding 向量
+    """
+    # 设置 AWS 环境变量（优先从参数取）
+    if aws_access_key_id:
+        os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
+    if aws_secret_access_key:
+        os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
+    if aws_session_token:
+        os.environ["AWS_SESSION_TOKEN"] = aws_session_token
+
+    bedrock_client = boto3.client("bedrock-sage.kernels.runtime")
+
+    model_provider = model.split(".")[0]
+
+    if model_provider == "amazon":
+        if "v2" in model:
+            body = json.dumps({
+                "inputText": text,
+                "embeddingTypes": ["float"],
+            })
+        elif "v1" in model:
+            body = json.dumps({"inputText": text})
+        else:
+            raise ValueError(f"Model {model} is not supported!")
+
+        response = bedrock_client.invoke_model(
+            modelId=model,
+            body=body,
+            accept="application/json",
+            contentType="application/json",
+        )
+        response_body = json.loads(response["body"].read())
+        return response_body["embedding"]
+
+    elif model_provider == "cohere":
+        body = json.dumps({
+            "texts": [text],
+            "input_type": "search_document",
+            "truncate": "NONE",
+        })
+
+        response = bedrock_client.invoke_model(
+            modelId=model,
+            body=body,
+            accept="application/json",
+            contentType="application/json",
+        )
+        response_body = json.loads(response["body"].read())
+        return response_body["embeddings"][0]
+
+    else:
+        raise ValueError(f"Model provider '{model_provider}' is not supported!")

sage/middleware/utils/embedding/embedding_api.py

@@ -0,0 +1,12 @@
+from sage.middleware.utils.embedding.embedding_model import 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;
+    while name(method) = "openai",if you need call other APIs which are compatible with openai,set the params:base_url,api_key,model;
+    while name(method) = "jina/siliconcloud/cohere",please set the params:api_key,model;
+    Example:operator_test.py
+    """
+    return EmbeddingModel(method=name,**kwargs)

sage/middleware/utils/embedding/embedding_model.py

@@ -0,0 +1,150 @@
+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__), "../../../../..")))
+
+import time
+
+from dotenv import load_dotenv
+
+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
+
+
+class EmbeddingModel:
+    # def __init__(self, method: str = "openai", model: str = "mistral-embed",
+    #              base_url: str = None, api_key: str = None):
+    def __init__(self, method: str = "openai", **kwargs):
+        """
+        初始化 embedding table
+        :param method: 指定使用的 embedding 方法名称，例如 "openai" 或 "cohere" 或“hf"等
+        """
+        self.init_method = method
+        self.dim = None
+        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.method = method
+
+        # self.kwargs = {}
+        self.kwargs = kwargs
+        if method == "hf":
+            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")
+        elif method == "mockembedder":
+            # 初始化 mockembedder
+            self.kwargs["embed_model"] = mockembedder.MockTextEmbedder(
+                model_name="mock-model",
+                fixed_dim=kwargs.get("fixed_dim", 128)
+            )
+        self.embed_fn = self._get_embed_function(method)
+
+    def set_dim(self, model_name):
+        """
+        :param model_name:
+        :return:
+        """
+        dimension_mapping = {
+            "mistral_embed": 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
+        }
+        if model_name in dimension_mapping:
+            self.dim = dimension_mapping[model_name]
+        else:
+            raise ValueError(f"<UNK> embedding <UNK>{model_name}")
+
+    def get_dim(self):
+        return self.dim
+
+    def _get_embed_function(self, method: str):
+        """根据方法名返回对应的 embedding 函数"""
+        mapping = {
+            "openai": openai.openai_embed_sync,
+            "zhipu": zhipu.zhipu_embedding_sync,
+            "bedrock": bedrock.bedrock_embed_sync,
+            "hf": hf.hf_embed_sync,
+            "jina": jina.jina_embed_sync,
+            # "llama_index_impl": llama_index_impl.llama_index_embed,
+            "lollms": lollms.lollms_embed_sync,
+            "nvidia_openai": nvidia_openai.nvidia_openai_embed_sync,
+            "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(),
+            # "instructor": instructor.instructor_embed
+        }
+        if method not in mapping:
+            raise ValueError(f"不支持的 embedding 方法：{method}")
+
+        embed_fn = mapping[method]
+
+        return embed_fn
+
+    def _embed(self, text: str) -> list[float]:
+        """
+        异步执行 embedding 操作
+        :param text: 要 embedding 的文本
+        :param kwargs: embedding 方法可能需要的额外参数
+        :return: embedding 后的结果
+        """
+        return self.embed_fn(text, **self.kwargs)
+
+    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:
+    """
+    usage  参见sage/api/model/operator_test.py
+    while name(method) = "hf", please set the param:model;
+    while name(method) = "openai",if you need call other APIs which are compatible with openai,set the params:base_url,api_key,model;
+    while name(method) = "jina/siliconcloud/cohere",please set the params:api_key,model;
+    Example:operator_test.py
+    """
+    return EmbeddingModel(method=name,**kwargs)
+
+def main():
+    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__":
+    main()

sage/middleware/utils/embedding/hf.py

@@ -0,0 +1,90 @@
+import asyncio
+import copy
+import os
+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
+except ImportError:
+    raise ImportError(
+        "transformers package is required for HuggingFace embedding functionality. "
+        "Please install it via: pip install transformers"
+    )
+
+try:
+    import torch
+except ImportError:
+    raise ImportError(
+        "torch package is required for HuggingFace embedding functionality. "
+        "Please install it via: pip install torch"
+    )
+
+try:
+    import tenacity
+except ImportError:
+    raise ImportError(
+        "tenacity package is required for HuggingFace embedding functionality. "
+        "Please install it via: pip install tenacity"
+    )
+
+try:
+    import numpy
+except ImportError:
+    raise ImportError(
+        "numpy package is required for HuggingFace embedding functionality. "
+        "Please install it via: pip install numpy"
+    )
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+@lru_cache(maxsize=1)
+def initialize_hf_model(model_name):
+    hf_tokenizer = AutoTokenizer.from_pretrained(
+        model_name, device_map="auto", trust_remote_code=True
+    )
+    hf_model = AutoModelForCausalLM.from_pretrained(
+        model_name, device_map="auto", trust_remote_code=True
+    )
+    if hf_tokenizer.pad_token is None:
+        hf_tokenizer.pad_token = hf_tokenizer.eos_token
+
+    return hf_model, hf_tokenizer
+
+
+import torch
+
+def hf_embed_sync(text: str, tokenizer, embed_model) -> list[float]:
+    """
+    使用 HuggingFace 模型同步生成文本 embedding。
+
+    Args:
+        text (str): 输入文本
+        tokenizer: 已加载的 tokenizer
+        embed_model: 已加载的 PyTorch embedding 模型
+
+    Returns:
+        list[float]: embedding 向量
+    """
+    device = next(embed_model.parameters()).device
+    encoded_texts = tokenizer(
+        text, return_tensors="pt", padding=True, truncation=True
+    ).to(device)
+
+    with torch.no_grad():
+        outputs = embed_model(
+            input_ids=encoded_texts["input_ids"],
+            attention_mask=encoded_texts["attention_mask"],
+        )
+        embeddings = outputs.last_hidden_state.mean(dim=1)
+
+    if embeddings.dtype == torch.bfloat16:
+        return embeddings.detach().to(torch.float32).cpu()[0].tolist()
+    else:
+        return embeddings.detach().cpu()[0].tolist()
+
+

sage/middleware/utils/embedding/instructor.py

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

sage/middleware/utils/embedding/jina.py

@@ -0,0 +1,115 @@
+import asyncio
+import os
+import pipmaster as pm  # Pipmaster for dynamic library install
+
+# Dependencies should be installed via requirements.txt
+# tenacity is required for this module
+
+try:
+    import tenacity
+except ImportError:
+    raise ImportError(
+        "tenacity package is required for Jina embedding functionality. "
+        "Please install it via: pip install tenacity"
+    )
+
+try:
+    import aiohttp
+except ImportError:
+    raise ImportError(
+        "aiohttp package is required for Jina embedding functionality. "
+        "Please install it via: pip install aiohttp"
+    )
+
+import numpy as np
+
+
+async def fetch_data(url, headers, data):
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, headers=headers, json=data) as response:
+            response_json = await response.json()
+            data_list = response_json.get("data", [])
+            return data_list
+
+
+async def jina_embed(
+    text: str,
+    dimensions: int = 1024,
+    late_chunking: bool = False,
+    base_url: str = None,
+    api_key: str = None,
+    model:str="jina-embeddings-v3"
+) -> list[float]:
+    if api_key:
+        os.environ["JINA_API_KEY"] = api_key
+    url = "https://api.jina.ai/v1/embeddings" if not base_url else base_url
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {os.environ['JINA_API_KEY']}",
+    }
+    data = {
+        "model": f"{model}",
+        "normalized": True,
+        "embedding_type": "float",
+        "dimensions": f"{dimensions}",
+        "late_chunking": late_chunking,
+        "input": text,
+    }
+    data_list = await fetch_data(url, headers, data)
+    print(data_list)
+    return data_list[0]['embedding']
+
+import os
+import requests
+
+def jina_embed_sync(
+    text: str,
+    dimensions: int = 1024,
+    late_chunking: bool = False,
+    base_url: str = None,
+    api_key: str = None,
+    model: str = "jina-embeddings-v3"
+) -> list[float]:
+    """
+    同步版本：调用 Jina AI embedding API 获取嵌入向量
+
+    Args:
+        text: 待嵌入的文本
+        dimensions: 嵌入维度
+        late_chunking: 是否开启 late chunking
+        base_url: 自定义 API 地址（可选）
+        api_key: Jina API 密钥
+        model: 使用的模型名
+
+    Returns:
+        list[float]: 嵌入向量
+    """
+    if api_key:
+        os.environ["JINA_API_KEY"] = api_key
+
+    url = base_url or "https://api.jina.ai/v1/embeddings"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {os.environ['JINA_API_KEY']}",
+    }
+    payload = {
+        "model": model,
+        "normalized": True,
+        "embedding_type": "float",
+        "dimensions": dimensions,
+        "late_chunking": late_chunking,
+        "input": text,
+    }
+
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        data = response.json()
+        return data["data"][0]["embedding"]
+    except Exception as e:
+        raise RuntimeError(f"Jina API call failed: {str(e)}")
+
+
+
+
+

sage/middleware/utils/embedding/lollms.py

@@ -0,0 +1,100 @@
+import sys
+
+if sys.version_info < (3, 9):
+    from typing import AsyncIterator
+else:
+    from collections.abc import AsyncIterator
+import pipmaster as pm  # Pipmaster for dynamic library install
+
+# Dependencies should be installed via requirements.txt
+# aiohttp and tenacity are required for this module
+
+try:
+    import aiohttp
+except ImportError:
+    raise ImportError(
+        "aiohttp package is required for Lollms embedding functionality. "
+        "Please install it via: pip install aiohttp"
+    )
+
+try:
+    import tenacity
+except ImportError:
+    raise ImportError(
+        "tenacity package is required for Lollms embedding functionality. "
+        "Please install it via: pip install tenacity"
+    )
+
+
+
+async def lollms_embed(
+    text: str,
+    embed_model=None,
+    base_url="http://localhost:9600",
+    **kwargs,
+) -> list:
+    """
+    Generate embedding for a single text using lollms server.
+
+    Args:
+        text: The string to embed
+        embed_model: Model name (not used directly as lollms uses configured vectorizer)
+        base_url: URL of the lollms server
+        **kwargs: Additional arguments passed to the request
+
+    Returns:
+        list[float]: The embedding vector
+    """
+    api_key = kwargs.pop("api_key", None)
+    headers = (
+        {"Content-Type": "application/json", "Authorization": api_key}
+        if api_key
+        else {"Content-Type": "application/json"}
+    )
+
+    async with aiohttp.ClientSession(headers=headers) as session:
+        request_data = {"text": text}
+
+        async with session.post(
+            f"{base_url}/lollms_embed",
+            json=request_data,
+        ) as response:
+            result = await response.json()
+            return result["vector"]
+
+import requests
+
+def lollms_embed_sync(
+    text: str,
+    embed_model=None,
+    base_url="http://localhost:9600",
+    **kwargs,
+) -> list[float]:
+    """
+    同步版本：使用 lollms 本地服务生成 embedding。
+
+    Args:
+        text: 输入文本
+        embed_model: 模型名（未直接使用）
+        base_url: lollms 服务地址
+        **kwargs: 可选参数，例如 api_key
+
+    Returns:
+        list[float]: 生成的向量
+    """
+    api_key = kwargs.pop("api_key", None)
+    headers = {
+        "Content-Type": "application/json",
+    }
+    if api_key:
+        headers["Authorization"] = api_key
+
+    request_data = {"text": text}
+
+    try:
+        response = requests.post(f"{base_url}/lollms_embed", json=request_data, headers=headers)
+        response.raise_for_status()
+        result = response.json()
+        return result["vector"]
+    except Exception as e:
+        raise RuntimeError(f"lollms embedding request failed: {str(e)}")

sage/middleware/utils/embedding/mockembedder.py

@@ -0,0 +1,46 @@
+import torch
+import hashlib
+from typing import Optional
+
+class MockTextEmbedder:
+    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:
+        """生成固定维度的伪嵌入（相同文本输出相同）"""
+        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()
+        else:
+            return self._embed_with_sliding_window()
+
+    def _embed_single(self) -> torch.Tensor:
+        """模拟单文本嵌入"""
+        embedding = torch.randn(384)  # 模拟原始模型的中间维度
+        return self._adjust_dimension(embedding)
+
+    def _embed_with_sliding_window(self) -> torch.Tensor:
+        """模拟长文本滑动窗口嵌入"""
+        embeddings = torch.stack([torch.randn(384) for _ in range(3)])  # 模拟3个窗口
+        return self._adjust_dimension(embeddings.mean(dim=0))
+
+    def _adjust_dimension(self, embedding: torch.Tensor) -> torch.Tensor:
+        """保持与原代码一致的维度调整逻辑"""
+        if embedding.size(0) > 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