File size: 2,014 Bytes
03aa457 de800a4 03aa457 de800a4 03aa457 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from llama_index.llms.ollama import Ollama
from llama_index.core.llms import ChatMessage
import time
app = FastAPI()
class Generate(BaseModel):
response: str
duration: float
chat_history = []
model = Ollama(model="phi2", base_url="http://localhost:11434")
def generate_text(model: Ollama, prompt: str) -> {}:
if prompt == "":
return {
"response": "Please provide a prompt.",
"duration": str(0)
}
chat_history.append({
"role": "user",
"content": prompt
})
messages = [ChatMessage(role=msg["role"], content=msg["content"]) for msg in chat_history]
start_time = time.time()
response_gen = model.stream_chat(messages)
full_response = ""
try:
for response_chunk in response_gen:
full_response += response_chunk.delta
except Exception as e:
return {
"response": f"Error: {str(e)}",
"duration": 0
}
duration = time.time() - start_time
chat_history.append({
"role": "assistant",
"content": full_response
})
return {
"response": full_response,
"duration": f"{duration:.2f}"
}
@app.get("/")
async def root():
return {"message": "Hello World"}
@app.get("/health")
async def health_check():
try:
import requests
response = requests.get("http://localhost:11434/api/version")
ollama_status = "OK" if response.status_code == 200 else "Not available"
except:
ollama_status = "Error"
return {
"status": "healthy",
"ollama_status": ollama_status,
"models_loaded": model is not None
}
@app.post("/api/generate", summary="Generate text from prompt", tags=["Generate"], response_model=Generate)
def inference(input_prompt: str):
return generate_text(model, input_prompt)
|