Remove lightrag_api_open_webui_demo.py file
This commit is contained in:
@@ -1,140 +0,0 @@
|
||||
from datetime import datetime, timezone
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import StreamingResponse
|
||||
import inspect
|
||||
import json
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
|
||||
import os
|
||||
import logging
|
||||
from lightrag import LightRAG, QueryParam
|
||||
from lightrag.llm.ollama import ollama_model_complete, ollama_embed
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
|
||||
import nest_asyncio
|
||||
|
||||
WORKING_DIR = "./dickens"
|
||||
|
||||
logging.basicConfig(format="%(levelname)s:%(message)s", level=logging.INFO)
|
||||
|
||||
if not os.path.exists(WORKING_DIR):
|
||||
os.mkdir(WORKING_DIR)
|
||||
|
||||
rag = LightRAG(
|
||||
working_dir=WORKING_DIR,
|
||||
llm_model_func=ollama_model_complete,
|
||||
llm_model_name="qwen2.5:latest",
|
||||
llm_model_max_async=4,
|
||||
llm_model_max_token_size=32768,
|
||||
llm_model_kwargs={"host": "http://localhost:11434", "options": {"num_ctx": 32768}},
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=1024,
|
||||
max_token_size=8192,
|
||||
func=lambda texts: ollama_embed(
|
||||
texts=texts, embed_model="bge-m3:latest", host="http://127.0.0.1:11434"
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
with open("./book.txt", "r", encoding="utf-8") as f:
|
||||
rag.insert(f.read())
|
||||
|
||||
# Apply nest_asyncio to solve event loop issues
|
||||
nest_asyncio.apply()
|
||||
|
||||
app = FastAPI(title="LightRAG", description="LightRAG API open-webui")
|
||||
|
||||
|
||||
# Data models
|
||||
MODEL_NAME = "LightRAG:latest"
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
role: Optional[str] = None
|
||||
content: str
|
||||
|
||||
|
||||
class OpenWebUIRequest(BaseModel):
|
||||
stream: Optional[bool] = None
|
||||
model: Optional[str] = None
|
||||
messages: list[Message]
|
||||
|
||||
|
||||
# API routes
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def index():
|
||||
return "Set Ollama link to http://ip:port/ollama in Open-WebUI Settings"
|
||||
|
||||
|
||||
@app.get("/ollama/api/version")
|
||||
async def ollama_version():
|
||||
return {"version": "0.4.7"}
|
||||
|
||||
|
||||
@app.get("/ollama/api/tags")
|
||||
async def ollama_tags():
|
||||
return {
|
||||
"models": [
|
||||
{
|
||||
"name": MODEL_NAME,
|
||||
"model": MODEL_NAME,
|
||||
"modified_at": "2024-11-12T20:22:37.561463923+08:00",
|
||||
"size": 4683087332,
|
||||
"digest": "845dbda0ea48ed749caafd9e6037047aa19acfcfd82e704d7ca97d631a0b697e",
|
||||
"details": {
|
||||
"parent_model": "",
|
||||
"format": "gguf",
|
||||
"family": "qwen2",
|
||||
"families": ["qwen2"],
|
||||
"parameter_size": "7.6B",
|
||||
"quantization_level": "Q4_K_M",
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@app.post("/ollama/api/chat")
|
||||
async def ollama_chat(request: OpenWebUIRequest):
|
||||
resp = rag.query(
|
||||
request.messages[-1].content, param=QueryParam(mode="hybrid", stream=True)
|
||||
)
|
||||
if inspect.isasyncgen(resp):
|
||||
|
||||
async def ollama_resp(chunks):
|
||||
async for chunk in chunks:
|
||||
yield (
|
||||
json.dumps(
|
||||
{
|
||||
"model": MODEL_NAME,
|
||||
"created_at": datetime.now(timezone.utc).strftime(
|
||||
"%Y-%m-%dT%H:%M:%S.%fZ"
|
||||
),
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": chunk,
|
||||
},
|
||||
"done": False,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
+ b"\n"
|
||||
) # the b"\n" is important
|
||||
|
||||
return StreamingResponse(ollama_resp(resp), media_type="application/json")
|
||||
else:
|
||||
return resp
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
return {"status": "healthy"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host="0.0.0.0", port=8020)
|
Reference in New Issue
Block a user