feat: rag

openai/openai_client_a.py

@@ -0,0 +1,136 @@
+
+from typing import Optional,Iterator
+import os
+import requests
+import json
+class Message:
+    def __init__(self,data):
+        self.role = data.get('role'),
+        self.content = data.get('content')
+class Choice:
+    def __init__(self,data):
+        self.index = data.get('index')
+        self.message = Message(data.get('message',{}))
+        self.finish_reason = data.get('finish_reason')
+class ChatCompletionResponse:
+    def __init__(self,data):
+        self.id = data.get('id')
+        self.object = data.get('object')
+        self.created = data.get('created')
+        self.model = data.get('model')
+        choices_data = data.get('choices',[])
+        self.choices = [
+            Choice(choice_data) for choice_data in choices_data
+        ]
+        usage_data = data.get('usage',{})
+        self.usage = {
+            "prompt_tokens":usage_data.get("prompt_tokens"),
+            "completion_tokens":usage_data.get("completion_tokens"),
+            "total_tokens":usage_data.get("total_tokens"),
+        }
+class DeltaMessage:
+    def __init__(self,data):
+        self.content = data.get('content')  
+        self.role = data.get('role')        
+class DeltaChoice:
+    def __init__(self,data):
+        self.index = data.get('index')  
+        self.delta = DeltaMessage(data.get('delta',{}) ) 
+        self.finish_reason = data.get('finish_reason')       
+#流式响应数据块，表示流式响应中的一个数据块        
+class StreamChunk:
+    def __init__(self,data):
+        self.id = data.get('id')
+        self.object = data.get('object')
+        self.created = data.get('created')
+        self.model = data.get('model')
+        choices_data = data.get('choices',[])
+        self.choices = [DeltaChoice(choice_data) for choice_data in choices_data]
+
+class Stream:
+    def __init__(self,response:requests.Response):
+        self.response=response
+    def __enter__(self):
+        return self
+    def __exit__(self,exc_type,exc_val,exc_tb):
+        self.response.close()
+    def __iter__(self)->Iterator[StreamChunk]:
+        #迭代器方法，逐个返回流式数据块
+        try:
+            # 逐行读取响应的内容(SSE格式)
+            for line in self.response.iter_lines(decode_unicode=True):
+                #print(line)
+                if not line.strip():
+                    continue
+                if line.startswith('data: '):
+                    json_str = line[6:]
+                    # 如果遇到[DONE]说明流式输出结束
+                    if json_str.strip()=="[DONE]":
+                        break
+                    try:
+                        data = json.loads(json_str)
+                        yield StreamChunk(data)
+                    except json.JSONDecodeError:
+                        continue    
+
+        finally:
+            self.response.close()    
+
+class ChatCompletions:
+    def __init__(self,client):
+        self._client = client
+    def create(self,model,messages,max_tokens=1024,temperature=0.7,stream:bool=False,**kwargs):
+        url = f"{self._client.base_url}/chat/completions"
+        body = {
+            "model":model,
+            "messages":messages
+        }
+        if max_tokens is not None:
+            body["max_tokens"]=max_tokens
+        if temperature is not None:
+            body["temperature"]=temperature  
+        if stream:
+            body["stream"]=True    
+        #添加额外的参数到请求体中
+        body.update(kwargs)   
+        headers = {
+            "Authorization":f"Bearer {self._client.api_key}",
+            "Content-Type":"application/json"
+        }   
+        if stream:
+            response = requests.post(
+                url,
+                headers=headers,
+                json=body,
+                timeout=self._client.timeout,
+                stream=True#告诉openai的服务器我要使用流式输出
+            )
+            response.raise_for_status()
+            return Stream(response)
+        else:
+            response = requests.post(
+                url,
+                headers=headers,
+                json=body,
+                timeout=self._client.timeout
+            )
+            # 如果响应的状态不是2XX的话，主抛异常
+            response.raise_for_status()
+            return ChatCompletionResponse(response.json())
+
+class ChatResource:
+    def __init__(self,client):
+        self._client = client
+    @property
+    def completions(self)->ChatCompletions:
+        return ChatCompletions(self._client)
+class OpenAI:
+    def __init__(self,api_key:Optional[str]=None,base_url:str="https://api.openai.com/v1",timeout:float=60.0):
+        self.api_key=api_key or os.getenv("OPENAI_API_KEY")
+        if not self.api_key:
+            raise ValueError(f"API密钥未设置，请设置api_key参数或者环境变量OPENAI_API_KEY")
+        self.base_url = base_url.rstrip('/')
+        self.timeout = timeout
+    @property    
+    def chat(self)->ChatResource:
+        return ChatResource(self)