Mercurial

comparison mrjunejune/inference/litellm_proxy.py @ 260:1f9877b637e9

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