PyxiLabs commited on
Commit
310599e
·
verified ·
1 Parent(s): 16aa2f8

Create ollama.py

Browse files
Files changed (1) hide show
  1. ollama.py +72 -0
ollama.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, Request, HTTPException
2
+ from fastapi.responses import StreamingResponse
3
+ import httpx
4
+ import logging
5
+
6
+ app = FastAPI()
7
+ OLLAMA_BASE_URL = "http://localhost:11434"
8
+ client = httpx.AsyncClient(base_url=OLLAMA_BASE_URL)
9
+
10
+ # Set up logging
11
+ logging.basicConfig(level=logging.INFO)
12
+ logger = logging.getLogger("uvicorn")
13
+
14
+ @app.middleware("http")
15
+ async def log_requests(request: Request, call_next):
16
+ logger.info(f"Request: {request.method} {request.url}")
17
+ response = await call_next(request)
18
+ return response
19
+
20
+ @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
21
+ async def reverse_proxy(request: Request, path: str):
22
+ try:
23
+ # Build the Ollama URL
24
+ url = f"{OLLAMA_BASE_URL}/{path}"
25
+
26
+ # Forward headers excluding those that might cause issues
27
+ headers = {
28
+ key: value
29
+ for key, value in request.headers.items()
30
+ if key.lower() not in ["host", "content-length"]
31
+ }
32
+
33
+ # Forward the request to Ollama
34
+ req = client.build_request(
35
+ method=request.method,
36
+ url=url,
37
+ headers=headers,
38
+ content=await request.body(),
39
+ params=request.query_params
40
+ )
41
+
42
+ response = await client.send(req, stream=True)
43
+
44
+ # Handle streaming responses
45
+ if "text/event-stream" in response.headers.get("content-type", ""):
46
+ return StreamingResponse(
47
+ response.aiter_bytes(),
48
+ media_type=response.headers.get("content-type"),
49
+ headers=dict(response.headers)
50
+ )
51
+
52
+ return StreamingResponse(
53
+ response.aiter_bytes(),
54
+ media_type=response.headers.get("content-type"),
55
+ headers=dict(response.headers)
56
+ )
57
+
58
+ except httpx.ConnectError:
59
+ raise HTTPException(
60
+ status_code=503,
61
+ detail="Ollama server unavailable"
62
+ )
63
+ except Exception as e:
64
+ logger.error(f"Proxy error: {str(e)}")
65
+ raise HTTPException(
66
+ status_code=500,
67
+ detail=f"Proxy error: {str(e)}"
68
+ )
69
+
70
+ if __name__ == "__main__":
71
+ import uvicorn
72
+ uvicorn.run(app, host="0.0.0.0", port=7860)