isage-middleware 0.1.3.1__py3-none-any.whl → 0.2.0.0__py3-none-any.whl

sage/middleware/utils/embedding/lollms.py

@@ -1,100 +0,0 @@
-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

@@ -1,46 +0,0 @@
-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

sage/middleware/utils/embedding/nvidia_openai.py

@@ -1,97 +0,0 @@
-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

@@ -1,97 +0,0 @@
-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

@@ -1,112 +0,0 @@
-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

@@ -1,133 +0,0 @@
-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

@@ -1,85 +0,0 @@
-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)}")
-