From 4df56571bec2c6f36e661bbd210f6dd907d5add9 Mon Sep 17 00:00:00 2001
From: Shroominic <db2002dominic@gmail.com>
Date: Sat, 1 Nov 2025 13:33:44 +0800
Subject: [PATCH] fix typing + linting issues

---
 routstr/upstreams/ollama.py   | 15 +++++++++------
 routstr/upstreams/upstream.py |  6 +-----
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/routstr/upstreams/ollama.py b/routstr/upstreams/ollama.py
index 3fb7dbda..e999e221 100644
--- a/routstr/upstreams/ollama.py
+++ b/routstr/upstreams/ollama.py
@@ -3,10 +3,13 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 import httpx
+from fastapi import Request
+from fastapi.responses import Response, StreamingResponse
 
 from .upstream import UpstreamProvider
 
 if TYPE_CHECKING:
+    from ..core.db import ApiKey, AsyncSession
     from ..payment.models import Model
 
 from ..core.logging import get_logger
@@ -43,15 +46,15 @@ class OllamaUpstreamProvider(UpstreamProvider):
 
     async def forward_request(
         self,
-        request,
+        request: Request,
         path: str,
         headers: dict,
-        request_body: bytes,
-        key,
+        request_body: bytes | None,
+        key: ApiKey,
         max_cost_for_model: int,
-        session,
-        model_obj,
-    ):
+        session: AsyncSession,
+        model_obj: Model,
+    ) -> Response | StreamingResponse:
         """Override to use OpenAI-compatible endpoint for proxy requests."""
         if path.startswith("v1/"):
             path = path.replace("v1/", "")
diff --git a/routstr/upstreams/upstream.py b/routstr/upstreams/upstream.py
index cf54558f..9d3b1e35 100644
--- a/routstr/upstreams/upstream.py
+++ b/routstr/upstreams/upstream.py
@@ -4,13 +4,9 @@ import json
 import re
 import traceback
 from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING, Mapping
+from typing import Mapping
 
 import httpx
-
-if TYPE_CHECKING:
-    from ..payment.cost_caculation import CostData, MaxCostData
-
 from fastapi import BackgroundTasks, HTTPException, Request
 from fastapi.responses import Response, StreamingResponse