Merge pull request #7696 from ishiland/main

feat: add ollama ps endpoint
This commit is contained in:
Timothy Jaeryang Baek 2024-12-09 16:12:07 -08:00 committed by GitHub
commit a4ed027498
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -432,6 +432,25 @@ async def get_ollama_versions(url_idx: Optional[int] = None):
else:
return {"version": False}
@app.get("/api/ps")
async def get_ollama_loaded_models(user=Depends(get_verified_user)):
"""
List models that are currently loaded into Ollama memory, and which node they are loaded on.
"""
if app.state.config.ENABLE_OLLAMA_API:
tasks = [
aiohttp_get(
f"{url}/api/ps",
app.state.config.OLLAMA_API_CONFIGS.get(url, {}).get("key", None),
)
for url in app.state.config.OLLAMA_BASE_URLS
]
responses = await asyncio.gather(*tasks)
return dict(zip(app.state.config.OLLAMA_BASE_URLS, responses))
else:
return {}
class ModelNameForm(BaseModel):
name: str