Merge pull request #3549 from open-webui/dev

refac: ollama non streaming response
This commit is contained in:
Timothy Jaeryang Baek 2024-06-30 01:40:26 -07:00 committed by GitHub
commit 7bc88eb00d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 12 deletions

View File

@ -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( if stream:
r.content, return StreamingResponse(
status_code=r.status, r.content,
headers=dict(r.headers), status_code=r.status,
background=BackgroundTask(cleanup_response, response=r, session=session), 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")

View File

@ -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