首次提交：上传本地文件夹

utils/llm.py

@@ -0,0 +1,124 @@
+# import requests
+
+
+# class LLM:
+#     def __init__(self, model_uid, server_url):
+
+#         self.model_uid = model_uid
+#         self.server_url = server_url.rstrip("/")
+
+#     def generate(self, prompt, max_tokens=32000, temperature=0.2, **kwargs):
+
+#         url = f"{self.server_url}/v1/completions"
+#         headers = {"Content-Type": "application/json"}
+#         data = {
+#             "model": self.model_uid,
+#             "prompt": prompt,
+#             "max_tokens": max_tokens,
+#             "temperature": temperature,
+#             **kwargs,
+#         }
+
+#         try:
+#             response = requests.post(url, headers=headers, json=data)
+#             response.raise_for_status()
+#             result = response.json()
+#             return result["choices"][0]["text"]
+#         except requests.exceptions.RequestException as e:
+#             raise RuntimeError(f"请求失败: {e}")
+#         except KeyError:
+#             raise ValueError("响应格式错误，无法解析生成结果")
+
+
+# llm = LLM(model_uid="QwQ-32b", server_url="http://172.20.0.145:9995")
+
+
+class Embedding:
+    def __init__(self, model_uid, server_url):
+
+        self.model_uid = model_uid
+        self.server_url = server_url.rstrip("/")
+
+    def embed(self, text):
+
+        url = f"{self.server_url}/v1/embeddings"
+        headers = {"Content-Type": "application/json"}
+        data = {"model": self.model_uid, "input": text}
+
+        try:
+            response = requests.post(url, headers=headers, json=data)
+            response.raise_for_status()
+            result = response.json()
+            return result["data"][0]["embedding"]
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"请求失败: {e}")
+        except KeyError:
+            raise ValueError("响应格式错误，无法解析嵌入结果")
+
+
+embedding = Embedding(model_uid="bge-m3", server_url="http://10.1.16.39:9995")
+
+from langchain_openai import ChatOpenAI
+import requests
+from typing import Any, Dict, List, Optional
+from langchain_core.language_models import BaseLLM
+from langchain_core.callbacks import CallbackManagerForLLMRun
+from langchain_core.outputs import LLMResult, Generation
+from langchain_community.llms.yi import YiLLM
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+class SiliconFlowLLM(BaseLLM):
+    """自定义硅基流动大模型调用类"""
+
+    api_url: str
+    api_key: str
+    model: str
+
+    def _generate(
+        self,
+        prompts: List[str],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> LLMResult:
+        from langchain_core.outputs import Generation, LLMResult
+
+        generations = []
+        for prompt in prompts:
+            try:
+                headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+                payload = {
+                    "model": self.model,
+                    "messages": [{"role": "user", "content": prompt}],
+                }
+                response = requests.post(self.api_url, json=payload, headers=headers)
+                response.raise_for_status()
+                text = response.json()["choices"][0]["message"]["content"]
+                generations.append([Generation(text=text)])
+            except Exception as e:
+                raise Exception(f"调用硅基流动API失败: {str(e)}")
+
+        return LLMResult(generations=generations)
+
+    @property
+    def _llm_type(self) -> str:
+        return "siliconflow"
+
+
+search_llm = SiliconFlowLLM(
+    api_url="https://api.siliconflow.cn/v1/chat/completions",
+    api_key="sk-bbeamiumkouptsrueilgufqqyuumelcsivxwjbdugqwsqhwj",
+    model="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
+)
+
+keyword_llm = SiliconFlowLLM(
+    api_url="https://api.siliconflow.cn/v1/chat/completions",
+    api_key="sk-bbeamiumkouptsrueilgufqqyuumelcsivxwjbdugqwsqhwj",
+    model="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
+)
+
+rewrite_llm = SiliconFlowLLM(
+    api_url="https://api.siliconflow.cn/v1/chat/completions",
+    api_key="sk-bbeamiumkouptsrueilgufqqyuumelcsivxwjbdugqwsqhwj",
+    model="Qwen/Qwen2.5-72B-Instruct",
+)