Mercurial

annotate mrjunejune/inference/litellm_proxy.py @ 279:b3b547563ec7

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Add Google connector service and agent wiki Implement the C/Seobeo Google Drive and Gmail connector with encrypted OAuth storage, Zenbu authentication, browser testing, AI tool discovery, chunked HTTP decoding, and Bazel coverage. Consolidate repository guidance into progressive wiki documentation and enforce arena-first allocation for new first-party C code. Co-authored-by: Copilot <[email protected]> Copilot-Session: 84c338fd-0939-4bb3-b7f3-1062eb213e5d
author MrJuneJune <me@mrjunejune.com>
date Mon, 17 Aug 2026 22:22:36 -0700
parents 1f9877b637e9
children
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
260
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
1 from __future__ import annotations
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
2
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
3 import argparse
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
4 import asyncio
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
5 import json
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
6 import os
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
7 import pathlib
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
8 from collections.abc import AsyncIterator, Awaitable, Callable
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
9 from typing import Any
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
10
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
11 import litellm
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
12 import uvicorn
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
13 import yaml
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
14 from fastapi import FastAPI, HTTPException, Request
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
15 from fastapi.responses import JSONResponse, StreamingResponse
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
16
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
17 Completion = Callable[..., Awaitable[Any]]
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
18
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
19
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
20 def resolve_token_directory(explicit: str | None = None) -> pathlib.Path:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
21 if explicit:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
22 return pathlib.Path(explicit).expanduser().resolve()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
23 configured = os.environ.get("GITHUB_COPILOT_TOKEN_DIR")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
24 if configured:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
25 return pathlib.Path(configured).expanduser().resolve()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
26 state_root = os.environ.get("MRJUNEJUNE_INFERENCE_STATE")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
27 if state_root:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
28 return (
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
29 pathlib.Path(state_root).expanduser().resolve()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
30 / "litellm-copilot"
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
31 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
32 xdg_state = os.environ.get("XDG_STATE_HOME")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
33 root = (
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
34 pathlib.Path(xdg_state).expanduser()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
35 if xdg_state
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
36 else pathlib.Path.home() / ".local" / "state"
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
37 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
38 return (root / "mrjunejune" / "inference" / "litellm-copilot").resolve()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
39
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
40
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
41 def _serialize(value: Any) -> dict[str, Any]:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
42 if hasattr(value, "model_dump"):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
43 return value.model_dump(exclude_none=True)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
44 if hasattr(value, "dict"):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
45 return value.dict()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
46 if isinstance(value, dict):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
47 return value
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
48 raise TypeError("LiteLLM returned an unsupported response")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
49
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
50
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
51 def create_app(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
52 completion: Completion = litellm.acompletion,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
53 *,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
54 model_alias: str = "jrpg-copilot",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
55 upstream_model: str = "github_copilot/gpt-4",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
56 ) -> FastAPI:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
57 app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
58
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
59 @app.get("/health/liveliness")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
60 async def health() -> dict[str, str]:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
61 return {"status": "ready"}
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
62
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
63 @app.get("/health/readiness")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
64 async def readiness() -> dict[str, str]:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
65 from litellm.llms.github_copilot.authenticator import Authenticator
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
66
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
67 try:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
68 await asyncio.to_thread(Authenticator().get_api_key)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
69 except Exception as error:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
70 raise HTTPException(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
71 status_code=503,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
72 detail="GitHub Copilot authentication unavailable",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
73 ) from error
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
74 return {"status": "ready"}
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
75
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
76 @app.post("/v1/chat/completions")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
77 async def chat(request: Request):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
78 master_key = os.environ.get("LITELLM_MASTER_KEY", "")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
79 authorization = request.headers.get("Authorization", "")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
80 if not master_key or authorization != f"Bearer {master_key}":
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
81 raise HTTPException(status_code=401, detail="Invalid gateway key")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
82
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
83 payload = await request.json()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
84 if not isinstance(payload, dict) or payload.get("model") != model_alias:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
85 raise HTTPException(status_code=400, detail="Unknown model alias")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
86 payload = dict(payload)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
87 payload["model"] = upstream_model
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
88 configured_max = int(os.environ.get("LITELLM_MAX_OUTPUT_TOKENS", "1024"))
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
89 requested_max = payload.get("max_tokens")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
90 if not isinstance(requested_max, int) or requested_max > configured_max:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
91 payload["max_tokens"] = configured_max
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
92
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
93 try:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
94 response = await completion(**payload)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
95 except Exception as error:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
96 raise HTTPException(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
97 status_code=502,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
98 detail="GitHub Copilot provider request failed",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
99 ) from error
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
100
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
101 if payload.get("stream"):
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
102 async def events() -> AsyncIterator[str]:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
103 async for chunk in response:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
104 yield (
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
105 "data: "
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
106 + json.dumps(_serialize(chunk), separators=(",", ":"))
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
107 + "\n\n"
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
108 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
109 yield "data: [DONE]\n\n"
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
110
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
111 return StreamingResponse(events(), media_type="text/event-stream")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
112 return JSONResponse(_serialize(response))
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
113
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
114 return app
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
115
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
116
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
117 def _load_model(config_path: str) -> tuple[str, str]:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
118 with open(config_path, encoding="utf-8") as config_file:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
119 config = yaml.safe_load(config_file)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
120 models = config.get("model_list", [])
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
121 if len(models) != 1:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
122 raise ValueError("LiteLLM config must contain exactly one model")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
123 model = models[0]
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
124 return model["model_name"], model["litellm_params"]["model"]
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
125
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
126
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
127 def main() -> None:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
128 parser = argparse.ArgumentParser(description="Minimal LiteLLM gateway")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
129 parser.add_argument("--host", default="127.0.0.1")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
130 parser.add_argument("--port", type=int, default=4000)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
131 parser.add_argument("--config")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
132 parser.add_argument("--authenticate", action="store_true")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
133 parser.add_argument(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
134 "--token-dir",
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
135 help=(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
136 "Persistent GitHub Copilot token directory. Defaults to "
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
137 "GITHUB_COPILOT_TOKEN_DIR, MRJUNEJUNE_INFERENCE_STATE, or "
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
138 "~/.local/state/mrjunejune/inference/litellm-copilot."
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
139 ),
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
140 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
141 args = parser.parse_args()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
142 if args.authenticate:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
143 from litellm.llms.github_copilot.authenticator import Authenticator
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
144
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
145 token_directory = resolve_token_directory(args.token_dir)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
146 token_directory.mkdir(parents=True, exist_ok=True, mode=0o700)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
147 token_directory.chmod(0o700)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
148 os.environ["GITHUB_COPILOT_TOKEN_DIR"] = str(token_directory)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
149 authenticator = Authenticator()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
150 authenticator.get_access_token()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
151 authenticator.get_api_key()
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
152 print(f"GitHub Copilot authentication stored in {token_directory}")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
153 return
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
154 if not args.config:
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
155 parser.error("--config is required unless --authenticate is used")
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
156 model_alias, upstream_model = _load_model(args.config)
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
157 uvicorn.run(
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
158 create_app(model_alias=model_alias, upstream_model=upstream_model),
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
159 host=args.host,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
160 port=args.port,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
161 access_log=False,
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
162 )
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
163
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
164
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
165 if __name__ == "__main__":
1f9877b637e9 Add Copilot-powered cyberpunk JRPG chat
MrJuneJune <mrjunejune@users.noreply.github.com>
parents:
diff changeset
166 main()