Przeglądaj źródła

Merge pull request #3549 from open-webui/dev

refac: ollama non streaming response
Timothy Jaeryang Baek 10 miesięcy temu
rodzic
commit
7bc88eb00d
2 zmienionych plików z 24 dodań i 12 usunięć
  1. 20 8
      backend/apps/ollama/main.py
  2. 4 4
      backend/requirements.txt

+ 20 - 8
backend/apps/ollama/main.py

@@ -153,7 +153,7 @@ async def cleanup_response(
         await session.close()
         await session.close()
 
 
 
 
-async def post_streaming_url(url: str, payload: str):
+async def post_streaming_url(url: str, payload: str, stream: bool = True):
     r = None
     r = None
     try:
     try:
         session = aiohttp.ClientSession(
         session = aiohttp.ClientSession(
@@ -162,12 +162,20 @@ async def post_streaming_url(url: str, payload: str):
         r = await session.post(url, data=payload)
         r = await session.post(url, data=payload)
         r.raise_for_status()
         r.raise_for_status()
 
 
-        return StreamingResponse(
-            r.content,
-            status_code=r.status,
-            headers=dict(r.headers),
-            background=BackgroundTask(cleanup_response, response=r, session=session),
-        )
+        if stream:
+            return StreamingResponse(
+                r.content,
+                status_code=r.status,
+                headers=dict(r.headers),
+                background=BackgroundTask(
+                    cleanup_response, response=r, session=session
+                ),
+            )
+        else:
+            res = await r.json()
+            await cleanup_response(r, session)
+            return res
+
     except Exception as e:
     except Exception as e:
         error_detail = "Open WebUI: Server Connection Error"
         error_detail = "Open WebUI: Server Connection Error"
         if r is not None:
         if r is not None:
@@ -963,7 +971,11 @@ async def generate_openai_chat_completion(
     url = app.state.config.OLLAMA_BASE_URLS[url_idx]
     url = app.state.config.OLLAMA_BASE_URLS[url_idx]
     log.info(f"url: {url}")
     log.info(f"url: {url}")
 
 
-    return await post_streaming_url(f"{url}/v1/chat/completions", json.dumps(payload))
+    return await post_streaming_url(
+        f"{url}/v1/chat/completions",
+        json.dumps(payload),
+        stream=payload.get("stream", False),
+    )
 
 
 
 
 @app.get("/v1/models")
 @app.get("/v1/models")

+ 4 - 4
backend/requirements.txt

@@ -6,7 +6,7 @@ python-multipart==0.0.9
 Flask==3.0.3
 Flask==3.0.3
 Flask-Cors==4.0.1
 Flask-Cors==4.0.1
 
 
-python-socketio==5.11.2
+python-socketio==5.11.3
 python-jose==3.3.0
 python-jose==3.3.0
 passlib[bcrypt]==1.7.4
 passlib[bcrypt]==1.7.4
 
 
@@ -32,10 +32,10 @@ google-generativeai==0.5.4
 
 
 langchain==0.2.0
 langchain==0.2.0
 langchain-community==0.2.0
 langchain-community==0.2.0
-langchain-chroma==0.1.1
+langchain-chroma==0.1.2
 
 
 fake-useragent==1.5.1
 fake-useragent==1.5.1
-chromadb==0.5.0
+chromadb==0.5.3
 sentence-transformers==2.7.0
 sentence-transformers==2.7.0
 pypdf==4.2.0
 pypdf==4.2.0
 docx2txt==0.8
 docx2txt==0.8
@@ -67,4 +67,4 @@ pytube==15.0.0
 
 
 extract_msg
 extract_msg
 pydub
 pydub
-duckduckgo-search~=6.1.5
+duckduckgo-search~=6.1.7