Update app.py

app.py

@@ -1,16 +1,13 @@
-
-
-
 import uvicorn
 import json
 import requests
 from flask import Flask, request, jsonify
-
+from flask import Response, stream_with_context
 
 
 app = Flask(__name__)
 
-rq = requests.sessions.Session()
+rq = requests.Session()
 
 model_names = [
         "meta-llama/Meta-Llama-3-70B-Instruct",
@@ -32,7 +29,7 @@ model_names = [
 
 
 
-def DeepinFra(Api:str, messages:list ,model:str, max_tokens: int = 512, temperature: float = 0.7):
+def DeepinFra_No_stream(Api:str, messages:list ,model:str = "meta-llama/Meta-Llama-3-70B-Instruct", max_tokens: int = 512, temperature: float = 0.7):
 
     url = "https://api.deepinfra.com/v1/openai/chat/completions"
     headers ={
@@ -62,6 +59,45 @@ def DeepinFra(Api:str, messages:list ,model:str, max_tokens: int = 512, temperat
         
         return "Response content: " + result.text
 
+def DeepinFra_stream(Api:str, messages:list ,model: str = "meta-llama/Meta-Llama-3-70B-Instruct", max_tokens: int = 512, temperature: float = 0.7):
+
+    url = "https://api.deepinfra.com/v1/openai/chat/completions"
+    headers ={
+        "Authorization" : f"Bearer {Api}",
+        'Content-Type': 'application/json',
+        'Accept': 'text/event-stream',
+    }
+
+
+
+
+    data = json.dumps(
+        {
+            'model': model,
+            'messages': messages,
+            'temperature': temperature,
+            'max_tokens': max_tokens,
+            'stream': True
+        }, separators=(',', ':')
+    )
+    
+    try:
+        result = rq.post(url=url, headers=headers, data=data, stream=True)
+        
+        for line in result.iter_lines():
+            if line:
+                line = line.decode('utf-8')
+                data_json = line.split('data: ')[1]
+                data = json.loads(data_json)
+                try:
+                    content = data['choices'][0]['delta']['content']
+                    yield content
+                except:
+                    break
+    except:
+
+        
+        return "Response content: " + result.text
 
 
 
@@ -70,13 +106,33 @@ def generate_text():
     data = request.json
     message = data.get("message")
     Api = data.get("api_key")
-    model_name = data.get("model_name", "mistralai/Mixtral-8x22B-Instruct-v0.1")
+    model_name = data.get("model_name", "meta-llama/Meta-Llama-3-70B-Instruct")
     max_tokens = data.get("max_tokens", 512)
     temperature = data.get("temperature", 0.7)
+    stream = data.get("stream", True)
 
     if not message or not Api:
         return jsonify({"error": "Missing required fields"}), 400
 
-    response = DeepinFra(Api=Api , messages=message, model=model_name, max_tokens=max_tokens, temperature=temperature)
+    def generate_response(stream: bool):
+        if stream:
+            for response in DeepinFra_stream(Api=Api, messages=message, model=model_name, max_tokens=max_tokens,
+                                              temperature=temperature):
+                yield json.dumps({"response": response}) + "\n"
+        else:
+            response = DeepinFra_No_stream(Api=Api, messages=message, model=model_name, max_tokens=max_tokens,
+                                           temperature=temperature)
+            yield json.dumps({"response": response}) + "\n"
+
+    return Response(stream_with_context(generate_response(stream)), content_type='application/json'), 200
+
+
+
+@app.route("/info", methods=["GET"])
+def get_info():
+    return jsonify({"model_names": model_names}), 200
+
+
 
-    return jsonify({"response": response}), 200
+if __name__=="__main__":
+    app.run()