isage-middleware 0.1.0__py3-none-any.whl

sage/middleware/utils/embedding/nvidia_openai.py

@@ -0,0 +1,97 @@
+import asyncio
+import sys
+import os
+
+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
+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",
+    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
+) -> list:
+    """
+    Generate embedding for a single text using NVIDIA NIM-compatible OpenAI API.
+
+    Returns:
+        list[float]: The embedding vector.
+    """
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
+
+    openai_async_client = (
+        AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url)
+    )
+
+    response = await openai_async_client.embeddings.create(
+        model=model,
+        input=text,
+        encoding_format=encode,
+        extra_body={"input_type": input_type, "truncate": trunc},
+    )
+
+    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
+) -> list[float]:
+    """
+    同步版本：使用 NVIDIA NIM 接口生成文本 embedding。
+
+    Args:
+        text: 输入文本
+        model: 使用的模型 ID
+        base_url: 接口地址（默认为 NVIDIA 接口）
+        api_key: API 密钥（使用 OPENAI_API_KEY 环境变量）
+        input_type: 输入类型（passage / query）
+        trunc: 截断策略
+        encode: 返回格式（float / base64）
+
+    Returns:
+        list[float]: 嵌入向量
+    """
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
+
+    client = OpenAI(base_url=base_url)
+
+    response = client.embeddings.create(
+        model=model,
+        input=text,
+        encoding_format=encode,
+        extra_body={
+            "input_type": input_type,
+            "truncate": trunc
+        },
+    )
+
+    return response.data[0].embedding

sage/middleware/utils/embedding/ollama.py

@@ -0,0 +1,97 @@
+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
+# ollama and tenacity are required for this module to work
+
+try:
+    import ollama
+except ImportError:
+    raise ImportError(
+        "ollama package is required for Ollama embedding functionality. "
+        "Please install it via: pip install ollama"
+    )
+
+try:
+    import tenacity
+except ImportError:
+    raise ImportError(
+        "tenacity package is required for Ollama embedding functionality. "
+        "Please install it via: pip install tenacity"
+    )
+
+
+
+import numpy as np
+from typing import Union
+
+
+
+
+async def ollama_embed(
+    text: str,
+    embed_model,
+    **kwargs
+) -> list:
+    """
+    Generate embedding for a single text using Ollama.
+
+    Args:
+        text: A single input string
+        embed_model: The name of the Ollama embedding model
+        **kwargs: Optional arguments (e.g. base_url, api_key)
+
+    Returns:
+        list[float]: The embedding vector
+    """
+    import ollama
+
+    api_key = kwargs.pop("api_key", None)
+    headers = {
+        "Content-Type": "application/json",
+        "User-Agent": "SAGE/0.0",
+    }
+    if api_key:
+        headers["Authorization"] = api_key
+    kwargs["headers"] = headers
+
+    ollama_client = ollama.Client(**kwargs)
+    data = ollama_client.embed(model=embed_model, input=text)
+    return data["embedding"]
+
+def ollama_embed_sync(
+    text: str,
+    embed_model,
+    **kwargs
+) -> list[float]:
+    """
+    同步版本：使用 Ollama 客户端生成 embedding 向量。
+
+    Args:
+        text: 输入文本
+        embed_model: 使用的模型名
+        **kwargs: 额外参数（可包含 base_url、api_key）
+
+    Returns:
+        list[float]: embedding 向量
+    """
+    import ollama
+
+    api_key = kwargs.pop("api_key", None)
+    headers = {
+        "Content-Type": "application/json",
+        "User-Agent": "SAGE/0.0",
+    }
+    if api_key:
+        headers["Authorization"] = api_key
+    kwargs["headers"] = headers
+
+    ollama_client = ollama.Client(**kwargs)
+    data = ollama_client.embed(model=embed_model, input=text)
+    return data["embedding"]

sage/middleware/utils/embedding/openai.py

@@ -0,0 +1,112 @@
+import asyncio
+import sys
+import os
+import logging
+
+
+
+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
+# openai is required for this module
+
+try:
+    from openai import AsyncOpenAI  # 确保导入了这个
+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",
+    base_url: str = None,
+    api_key: str = None,
+) -> list:
+    """
+    Generate embedding for a single text using OpenAI Embedding API.
+
+    Args:
+        text: Input string
+        model: OpenAI embedding model name
+        base_url: Optional custom endpoint
+        api_key: OpenAI API key
+
+    Returns:
+        list[float]: The embedding vector
+    """
+    if not api_key:
+        api_key = os.environ["OPENAI_API_KEY"]
+
+    default_headers = {
+        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) SAGE/0.0",
+        "Content-Type": "application/json",
+    }
+
+    openai_async_client = (
+        AsyncOpenAI(default_headers=default_headers, api_key=api_key)
+        if base_url is None
+        else AsyncOpenAI(
+            base_url=base_url, default_headers=default_headers, api_key=api_key
+        )
+    )
+
+    response = await openai_async_client.embeddings.create(
+        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",
+    base_url: str = None,
+    api_key: str = None,
+) -> list[float]:
+    """
+    同步生成 OpenAI embedding。
+
+    Args:
+        text: 输入文本
+        model: OpenAI embedding 模型名
+        base_url: 可选自定义 API endpoint
+        api_key: OpenAI API 密钥
+
+    Returns:
+        list[float]: embedding 向量
+    """
+    if not api_key:
+        api_key = os.environ["OPENAI_API_KEY"]
+
+    default_headers = {
+        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_8) SAGE/0.0",
+        "Content-Type": "application/json",
+    }
+
+    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
+        )
+    )
+
+    response: CreateEmbeddingResponse = openai_sync_client.embeddings.create(
+        model=model,
+        input=text,
+        encoding_format="float"
+    )
+
+    return response.data[0].embedding

sage/middleware/utils/embedding/siliconcloud.py

@@ -0,0 +1,133 @@
+import asyncio
+import os
+import sys
+
+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
+
+
+
+
+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,
+) -> list:
+    """
+    Generate embedding for a single text using SiliconCloud (NetEase Youdao).
+
+    Args:
+        text: Input string
+        model: Embedding model name
+        base_url: API endpoint
+        max_token_size: Max text length in tokens (cut if needed)
+        api_key: Your SiliconCloud API key
+
+    Returns:
+        list[float]: The embedding vector
+    """
+    if api_key and not api_key.startswith("Bearer "):
+        api_key = "Bearer " + api_key
+
+    headers = {
+        "Authorization": api_key,
+        "Content-Type": "application/json",
+    }
+
+    text = text[:max_token_size]
+    payload = {
+        "model": model,
+        "input": [text],
+        "encoding_format": "base64",
+    }
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(base_url, headers=headers, json=payload) as response:
+            content = await response.json()
+            if "code" in content:
+                raise ValueError(content)
+            base64_string = content["data"][0]["embedding"]
+
+    decode_bytes = base64.b64decode(base64_string)
+    n = len(decode_bytes) // 4
+    float_array = struct.unpack("<" + "f" * n, decode_bytes)
+    return list(float_array)
+
+import requests
+import base64
+import struct
+
+def siliconcloud_embedding_sync(
+    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[float]:
+    """
+    同步版本：使用 SiliconCloud (NetEase Youdao) 接口获取文本 embedding。
+
+    Args:
+        text: 输入文本
+        model: 模型名称
+        base_url: 接口地址
+        max_token_size: 截断长度（按字符）
+        api_key: API 密钥（可选，带或不带 "Bearer "）
+
+    Returns:
+        list[float]: embedding 向量
+    """
+    if api_key and not api_key.startswith("Bearer "):
+        api_key = "Bearer " + api_key
+
+    headers = {
+        "Authorization": api_key,
+        "Content-Type": "application/json",
+    }
+
+    text = text[:max_token_size]
+    payload = {
+        "model": model,
+        "input": [text],
+        "encoding_format": "base64",
+    }
+
+    try:
+        response = requests.post(base_url, headers=headers, json=payload)
+        response.raise_for_status()
+        content = response.json()
+
+        if "code" in content:
+            raise ValueError(f"SiliconCloud API error: {content}")
+
+        base64_string = content["data"][0]["embedding"]
+        decode_bytes = base64.b64decode(base64_string)
+        n = len(decode_bytes) // 4
+        float_array = struct.unpack("<" + "f" * n, decode_bytes)
+        return list(float_array)
+
+    except Exception as e:
+        raise RuntimeError(f"SiliconCloud embedding failed: {str(e)}")

sage/middleware/utils/embedding/zhipu.py

@@ -0,0 +1,85 @@
+import sys
+import re
+import json
+
+
+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
+# zhipuai is required for this module
+
+try:
+    import zhipuai
+except ImportError:
+    raise ImportError(
+        "zhipuai package is required for ZhipuAI embedding functionality. "
+        "Please install it via: pip install zhipuai"
+    )
+
+
+async def zhipu_embedding(
+    text: str,
+    model: str = "embedding-3",
+    api_key: str = None,
+    **kwargs
+) -> list:
+    """
+    Generate embedding for a single text using ZhipuAI.
+
+    Args:
+        text: Input string
+        model: Embedding model name
+        api_key: ZhipuAI API key
+        **kwargs: Additional arguments to ZhipuAI client
+
+    Returns:
+        list[float]: Embedding vector
+    """
+    try:
+        from zhipuai import ZhipuAI
+    except ImportError:
+        raise ImportError("Please install zhipuai before using this backend.")
+
+    client = ZhipuAI(api_key=api_key) if api_key else ZhipuAI()
+
+    try:
+        response = client.embeddings.create(model=model, input=[text], **kwargs)
+        return response.data[0].embedding
+    except Exception as e:
+        raise Exception(f"Error calling ChatGLM Embedding API: {str(e)}")
+    
+def zhipu_embedding_sync(
+    text: str,
+    model: str = "embedding-3",
+    api_key: str = None,
+    **kwargs
+) -> list[float]:
+    """
+    同步调用 ZhipuAI 生成 embedding 向量。
+
+    Args:
+        text: 输入字符串
+        model: 使用的 ZhipuAI 模型名称
+        api_key: API 密钥（可选）
+        **kwargs: 额外参数
+
+    Returns:
+        list[float]: 生成的 embedding 向量
+    """
+    try:
+        from zhipuai import ZhipuAI
+    except ImportError:
+        raise ImportError("Please install zhipuai before using this backend.")
+
+    client = ZhipuAI(api_key=api_key) if api_key else ZhipuAI()
+
+    try:
+        response = client.embeddings.create(model=model, input=[text], **kwargs)
+        return response.data[0].embedding
+    except Exception as e:
+        raise Exception(f"Error calling ChatGLM Embedding API: {str(e)}")
+