1from openai import OpenAI
 2
 3class DeepSeekChat:
 4    def __init__(self, api_key, base_url="https://api.deepseek.com"):
 5        self.client = OpenAI(api_key=api_key, base_url=base_url)
 6
 7    def chat(
 8        self,
 9        system_message,
10        user_message,
11        model="deepseek-chat",
12        max_tokens=1024,
13        temperature=0.7,
14        stream=True,
15    ):
16
17        response = self.client.chat.completions.create(
18            model=model,
19            messages=[
20                {"role": "system", "content": system_message},
21                {"role": "user", "content": user_message},
22            ],
23            max_tokens=max_tokens,
24            temperature=temperature,
25            stream=stream,
26        )
27
28        if stream:
29            return self._stream_response(response)
30        else:
31            return response.choices[0].message.content
32
33    def _stream_response(self, response):
34        full_response = ""
35        for chunk in response:
36            if chunk.choices[0].delta.content is not None:
37                content = chunk.choices[0].delta.content
38                print(content, end="", flush=True)
39                full_response += content
40        
41        print("\r\n===============我是分隔线===============")
42        return full_response
43
44# 使用示例
45if __name__ == "__main__":
46    deepseek_chat = DeepSeekChat(api_key="[你的 API Key]")
47    response = deepseek_chat.chat(
48        system_message="你是一个聪明的 AI 助手",
49        user_message="三国演义中战斗力排名前 10 的武将有谁？",
50        stream=True,
51    )
52    print("完整回答：", response)

1# 设置语言模型，使用 Ollama 提供的 Qwen2 7B 模型，并设置请求超时时间
2Settings.llm = Ollama(model="qwen2:7b", request_timeout=360.0)

 1from typing import Optional, List, Mapping, Any
 2
 3from llama_index.core import SimpleDirectoryReader, SummaryIndex
 4from llama_index.core.callbacks import CallbackManager
 5from llama_index.core.llms import (
 6    CustomLLM,
 7    CompletionResponse,
 8    CompletionResponseGen,
 9    LLMMetadata,
10)
11from llama_index.core.llms.callbacks import llm_completion_callback
12from llama_index.core import Settings
13
14class OurLLM(CustomLLM):
15    context_window: int = 3900
16    num_output: int = 256
17    model_name: str = "custom"
18    dummy_response: str = "My response"
19
20    @property
21    def metadata(self) -> LLMMetadata:
22        """Get LLM metadata."""
23        return LLMMetadata(
24            context_window=self.context_window,
25            num_output=self.num_output,
26            model_name=self.model_name,
27        )
28
29    @llm_completion_callback()
30    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
31        return CompletionResponse(text=self.dummy_response)
32
33    @llm_completion_callback()
34    def stream_complete(
35        self, prompt: str, **kwargs: Any
36    ) -> CompletionResponseGen:
37        response = ""
38        for token in self.dummy_response:
39            response += token
40            yield CompletionResponse(text=response, delta=token)
41
42# define our LLM
43Settings.llm = OurLLM()
44
45# define embed model
46Settings.embed_model = "local:BAAI/bge-base-en-v1.5"
47
48# Load the your data
49documents = SimpleDirectoryReader("./data").load_data()
50index = SummaryIndex.from_documents(documents)
51
52# Query and print response
53query_engine = index.as_query_engine()
54response = query_engine.query("<query_text>")
55print(response)

  1import os
  2import sys
  3import logging
  4from openai import OpenAI
  5from typing import Any, Generator
  6from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
  7from llama_index.embeddings.huggingface import HuggingFaceEmbedding
  8from llama_index.core.llms import (
  9    CustomLLM,
 10    CompletionResponse,
 11    CompletionResponseGen,
 12    LLMMetadata,
 13)
 14from llama_index.core.llms.callbacks import llm_completion_callback
 15from pydantic import BaseModel, Field
 16from dotenv import load_dotenv
 17from functools import cached_property
 18
 19# 配置日志 创建一个与当前模块同名的 logger
 20logging.basicConfig(level=logging.INFO)
 21logger = logging.getLogger(__name__)
 22
 23# 从环境变量获取 API 密钥
 24load_dotenv()
 25
 26API_KEY = os.getenv("DEEPSEEK_API_KEY")
 27if not API_KEY:
 28    raise ValueError("DEEPSEEK_API_KEY environment variable is not set")
 29
 30class DeepSeekChat(BaseModel):
 31    """DeepSeek 聊天模型的封装类。"""
 32
 33    api_key: str = Field(default=API_KEY)
 34    base_url: str = Field(default="https://api.deepseek.com")
 35
 36    class Config:
 37        """Pydantic 配置类。"""
 38
 39        arbitrary_types_allowed = True  # 允许模型接受任意类型的字段
 40        # 这增加了灵活性，但可能降低类型安全性
 41        # 在本类中，这可能用于允许使用 OpenAI 客户端等复杂类型
 42
 43    @cached_property
 44    def client(self) -> OpenAI:
 45        """创建并缓存 OpenAI 客户端实例。"""
 46        return OpenAI(api_key=self.api_key, base_url=self.base_url)
 47
 48    def chat(
 49        self,
 50        system_message: str,
 51        user_message: str,
 52        model: str = "deepseek-chat",
 53        max_tokens: int = 1024,
 54        temperature: float = 0.7,
 55        stream: bool = False,
 56    ) -> Any:
 57        """
 58        使用 DeepSeek API 发送聊天请求。
 59
 60        返回流式响应或完整响应内容。
 61        """
 62        try:
 63            response = self.client.chat.completions.create(
 64                model=model,
 65                messages=[
 66                    {"role": "system", "content": system_message},
 67                    {"role": "user", "content": user_message},
 68                ],
 69                max_tokens=max_tokens,
 70                temperature=temperature,
 71                stream=stream,
 72            )
 73            return response if stream else response.choices[0].message.content
 74        except Exception as e:
 75            logger.error(f"Error in DeepSeek API call: {e}")
 76            raise
 77
 78    def _stream_response(self, response) -> Generator[str, None, None]:
 79        """处理流式响应，逐块生成内容。"""
 80        for chunk in response:
 81            if chunk.choices[0].delta.content is not None:
 82                yield chunk.choices[0].delta.content
 83
 84class DeepSeekLLM(CustomLLM):
 85    """DeepSeek 语言模型的自定义实现。"""
 86
 87    deep_seek_chat: DeepSeekChat = Field(default_factory=DeepSeekChat)
 88
 89    @property
 90    def metadata(self) -> LLMMetadata:
 91        """返回 LLM 元数据。"""
 92        return LLMMetadata()
 93
 94    @llm_completion_callback()
 95    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
 96        """执行非流式完成请求。"""
 97        response = self.deep_seek_chat.chat(
 98            system_message="你是一个聪明的 AI 助手", user_message=prompt, stream=False
 99        )
100        return CompletionResponse(text=response)
101
102    @llm_completion_callback()
103    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
104        """执行流式完成请求。"""
105        response = self.deep_seek_chat.chat(
106            system_message="你是一个聪明的 AI 助手", user_message=prompt, stream=True
107        )
108
109        def response_generator():
110            """生成器函数，用于逐步生成响应内容。"""
111            response_content = ""
112            for chunk in self.deep_seek_chat._stream_response(response):
113                if chunk:
114                    response_content += chunk
115                    yield CompletionResponse(text=response_content, delta=chunk)
116
117        return response_generator()
118
119# 设置环境变量，禁用 tokenizers 的并行处理，以避免潜在的死锁问题
120os.environ["TOKENIZERS_PARALLELISM"] = "false"
121
122def main():
123    """主程序函数，演示如何使用 DeepSeekLLM 进行文档查询。"""
124    # 从指定目录加载文档数据
125    documents = SimpleDirectoryReader("data").load_data()
126
127    # 设置 LLM 和嵌入模型
128    Settings.llm = DeepSeekLLM()
129    Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-zh-v1.5")
130
131    # 创建索引和查询引擎
132    index = VectorStoreIndex.from_documents(documents)
133    query_engine = index.as_query_engine(streaming=True)
134
135    # 执行查询
136    print("查询结果：")
137    response = query_engine.query("作者学习过的编程语言有哪些？")
138
139    # 处理并输出响应
140    if hasattr(response, "response_gen"):
141        # 流式输出
142        for text in response.response_gen:
143            print(text, end="", flush=True)
144            sys.stdout.flush()  # 确保立即输出
145    else:
146        # 非流式输出
147        print(response.response, end="", flush=True)
148
149    print("\n 查询完成")
150
151if __name__ == "__main__":
152    main()