|
import json |
|
from fastapi import FastAPI, HTTPException |
|
import uvicorn |
|
import httpx |
|
from pydantic import BaseModel |
|
from typing import List, Union |
|
|
|
app = FastAPI() |
|
|
|
OLLAMA_URL = "http://localhost:11434" |
|
|
|
class EmbeddingRequest(BaseModel): |
|
input: Union[str, List[str]] |
|
model: str |
|
|
|
class EmbeddingResponse(BaseModel): |
|
object: str |
|
data: List[dict] |
|
model: str |
|
usage: dict |
|
|
|
@app.post("/v1/embeddings") |
|
async def create_embedding(request: EmbeddingRequest): |
|
async with httpx.AsyncClient() as client: |
|
if isinstance(request.input, str): |
|
request.input = [request.input] |
|
|
|
ollama_requests = [{"model": request.model, "prompt": text} for text in request.input] |
|
|
|
embeddings = [] |
|
|
|
|
|
for i, ollama_request in enumerate(ollama_requests): |
|
response = await client.post(f"{OLLAMA_URL}/api/embeddings", json=ollama_request) |
|
if response.status_code != 200: |
|
raise HTTPException(status_code=response.status_code, detail="Ollama API error") |
|
|
|
result = response.json() |
|
embeddings.append({ |
|
"object": "embedding", |
|
"embedding": result["embedding"], |
|
"index": i |
|
}) |
|
|
|
|
|
return EmbeddingResponse( |
|
object="list", |
|
data=embeddings, |
|
model=request.model, |
|
|
|
) |
|
|
|
if __name__ == "__main__": |
|
import argparse |
|
parser = argparse.ArgumentParser(description="Run the embedding proxy server") |
|
parser.add_argument("--port", type=int, default=11435, help="Port to run the server on") |
|
parser.add_argument("--host", type=str, default="http://localhost:11434", help="URL of the Ollama server") |
|
parser.add_argument("--reload", action="store_true", help="Enable auto-reload for development") |
|
args = parser.parse_args() |
|
|
|
OLLAMA_URL = args.host |
|
uvicorn.run("embedding_proxy:app", host="0.0.0.0", port=args.port, reload=args.reload) |
|
|
