enterprise-kb-rag/_chapter2_config.py

import hashlib
import math
import os
import re
from pathlib import Path
from typing import Iterable, List


CODE_DIR = Path(__file__).resolve().parent


def load_chapter2_env() -> None:
    try:
        from dotenv import load_dotenv
    except ModuleNotFoundError:
        return

    load_dotenv(CODE_DIR / ".env")
    load_dotenv()


load_chapter2_env()


def env_value(*names: str, default: str = "") -> str:
    for name in names:
        value = os.getenv(name)
        if value:
            return value
    return default


def openai_config() -> dict:
    return {
        "base_url": env_value("base_url", "BASE_URL", default="https://api.modelverse.cn/v1"),
        "api_key": env_value("api_key", "API_KEY", default=""),
        "model": env_value("model", "MODEL", default="gpt-5.5"),
    }


def get_chat_model(temperature: float = 0):
    provider = env_value("MODEL_PROVIDER", default="openai").lower()
    if provider == "ollama":
        from langchain_ollama import ChatOllama

        return ChatOllama(
            model=env_value("OLLAMA_MODEL", default="qwen2.5:7b"),
            base_url=env_value("OLLAMA_BASE_URL", default="http://127.0.0.1:11434"),
            temperature=temperature,
        )

    from langchain_openai import ChatOpenAI

    conf = openai_config()
    if not conf["api_key"]:
        raise RuntimeError("没有找到 api_key/API_KEY，请检查 代码/.env。")
    return ChatOpenAI(
        base_url=conf["base_url"],
        api_key=conf["api_key"],
        model=conf["model"],
        temperature=temperature,
    )


def get_ollama_llm(temperature: float = 0):
    from langchain_ollama import OllamaLLM

    return OllamaLLM(
        model=env_value("OLLAMA_MODEL", default="qwen2.5:7b"),
        base_url=env_value("OLLAMA_BASE_URL", default="http://127.0.0.1:11434"),
        temperature=temperature,
    )


def as_text(result) -> str:
    return getattr(result, "content", result)


class HashEmbeddings:
    """Small local embedding model for classroom demos without external services."""

    def __init__(self, dim: int = 64):
        self.dim = dim

    def _embed(self, text: str) -> List[float]:
        vector = [0.0] * self.dim
        tokens = re.findall(r"[\w\u4e00-\u9fff]+", text.lower())
        if not tokens:
            tokens = [text]
        for token in tokens:
            digest = hashlib.sha256(token.encode("utf-8")).digest()
            index = int.from_bytes(digest[:4], "big") % self.dim
            sign = 1.0 if digest[4] % 2 == 0 else -1.0
            vector[index] += sign
        norm = math.sqrt(sum(x * x for x in vector)) or 1.0
        return [x / norm for x in vector]

    def embed_query(self, text: str) -> List[float]:
        return self._embed(text)

    def embed_documents(self, texts: Iterable[str]) -> List[List[float]]:
        return [self._embed(text) for text in texts]


def get_embeddings():
    provider = env_value("EMBEDDING_PROVIDER", default="hash").lower()
    if provider == "ollama":
        from langchain_ollama import OllamaEmbeddings

        return OllamaEmbeddings(
            model=env_value("OLLAMA_EMBEDDING_MODEL", default="mxbai-embed-large"),
            base_url=env_value("OLLAMA_BASE_URL", default="http://127.0.0.1:11434"),
        )
    return HashEmbeddings()


def print_dependency_hint(error: ModuleNotFoundError) -> None:
    print(f"缺少依赖: {error.name}")
    print("可在 `代码` 目录执行: pip install -r requirements_2_tools.txt")