refac

main.py

@@ -5,7 +5,7 @@ from fastapi.concurrency import run_in_threadpool
 
 from starlette.responses import StreamingResponse, Response
 from pydantic import BaseModel, ConfigDict
-from typing import List, Union, Generator
+from typing import List, Union, Generator, Iterator
 
 
 import time
@@ -132,11 +132,19 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
                 res = get_response(
                     user_message,
                     messages=form_data.messages,
-                    body=form_data.model_dump_json(),
+                    body=form_data.model_dump(),
                 )
 
                 print(f"stream:true:{res}")
 
+                if isinstance(res, Iterator):
+                    for line in res:
+                        if line:
+                            # Decode the JSON data
+                            decoded_line = line.decode("utf-8")
+                            print(f"stream_content:Iterator:{decoded_line}")
+                            yield f"{decoded_line}\n\n"
+                else:
                     if isinstance(res, str):
                         message = stream_message_template(form_data.model, res)
                         print(f"stream_content:str:{message}")
@@ -171,10 +179,13 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
             res = get_response(
                 user_message,
                 messages=form_data.messages,
-                body=form_data.model_dump_json(),
+                body=form_data.model_dump(),
             )
             print(f"stream:false:{res}")
 
+            if isinstance(res, dict):
+                return res
+            else:
                 message = ""
 
                 if isinstance(res, str):

pipelines/examples/haystack_pipeline.py

@@ -80,7 +80,7 @@ class Pipeline:
 
     def get_response(
         self, user_message: str, messages: List[OpenAIChatMessage], body: dict
-    ) -> Union[str, Generator]:
+    ) -> Union[str, Generator, Iterator]:
         # This is where you can add your custom RAG pipeline.
         # Typically, you would retrieve relevant information from your knowledge base and synthesize it to generate a response.
 

pipelines/examples/llamaindex_ollama_github_pipeline.py

@@ -71,7 +71,7 @@ class Pipeline:
 
     def get_response(
         self, user_message: str, messages: List[OpenAIChatMessage], body: dict
-    ) -> Union[str, Generator]:
+    ) -> Union[str, Generator, Iterator]:
         # This is where you can add your custom RAG pipeline.
         # Typically, you would retrieve relevant information from your knowledge base and synthesize it to generate a response.
 

pipelines/examples/llamaindex_ollama_pipeline.py

@@ -31,7 +31,7 @@ class Pipeline:
 
     def get_response(
         self, user_message: str, messages: List[OpenAIChatMessage], body: dict
-    ) -> Union[str, Generator]:
+    ) -> Union[str, Generator, Iterator]:
         # This is where you can add your custom RAG pipeline.
         # Typically, you would retrieve relevant information from your knowledge base and synthesize it to generate a response.
 

pipelines/examples/llamaindex_pipeline.py

@@ -26,7 +26,7 @@ class Pipeline:
 
     def get_response(
         self, user_message: str, messages: List[OpenAIChatMessage], body: dict
-    ) -> Union[str, Generator]:
+    ) -> Union[str, Generator, Iterator]:
         # This is where you can add your custom RAG pipeline.
         # Typically, you would retrieve relevant information from your knowledge base and synthesize it to generate a response.
 

pipelines/examples/openai_pipeline.py

@@ -1,4 +1,4 @@
-from typing import List, Union, Generator
+from typing import List, Union, Generator, Iterator
 from schemas import OpenAIChatMessage
 import requests
 
@@ -19,31 +19,35 @@ class Pipeline:
 
     def get_response(
         self, user_message: str, messages: List[OpenAIChatMessage], body: dict
-    ) -> Union[str, Generator]:
+    ) -> Union[str, Generator, Iterator]:
         # This is where you can add your custom pipelines like RAG.'
         print(f"get_response:{__name__}")
 
         print(messages)
         print(user_message)
-        OPENAI_API_KEY = "your-api-key-here"
+        OPENAI_API_KEY = "your-openai-api-key-here"
 
         headers = {}
         headers["Authorization"] = f"Bearer {OPENAI_API_KEY}"
         headers["Content-Type"] = "application/json"
 
-        r = requests.request(
-            method="POST",
-            url="https://api.openai.com/v1",
-            data=body,
+        data = {**body, "model": "gpt-3.5-turbo"}
+
+        print(data)
+
+        try:
+            r = requests.post(
+                url="https://api.openai.com/v1/chat/completions",
+                json={**body, "model": "gpt-3.5-turbo"},
                 headers=headers,
                 stream=True,
             )
 
             r.raise_for_status()
 
-        # Check if response is SSE
-        if "text/event-stream" in r.headers.get("Content-Type", ""):
-            return r.iter_content(chunk_size=8192)
+            if data["stream"]:
+                return r.iter_lines()
             else:
-            response_data = r.json()
-            return f"{response_data['choices'][0]['text']}"
+                return r.json()
+        except Exception as e:
+            return f"Error: {e}"

pipelines/examples/pipeline_example.py

@@ -18,7 +18,7 @@ class Pipeline:
 
     def get_response(
         self, user_message: str, messages: List[OpenAIChatMessage], body: dict
-    ) -> Union[str, Generator]:
+    ) -> Union[str, Generator, Iterator]:
         # This is where you can add your custom pipelines like RAG.'
         print(f"get_response:{__name__}")