🛠️ fix: update Driver import paths in middleware files

- Updated Driver import paths in middleware files to resolve import issues.

Author: sudoskys

llmkira/middleware/llm_task.py

@@ -14,9 +14,9 @@
 from ..extra.user.schema import Cost
 from ..schema import RawMessage, Scraper
 from ..sdk.adapter import SCHEMA_GROUP, SingleModel
-from ..sdk.endpoint import Driver
 from ..sdk.endpoint.openai import Message, Openai, OpenaiResult
 from ..sdk.endpoint.schema import LlmRequest, LlmResult
+from ..sdk.endpoint.tee import Driver
 from ..sdk.memory.redis import RedisChatMessageHistory
 from ..sdk.schema import ToolCallCompletion, SystemMessage, Function, Tool
 from ..sdk.utils import sync

llmkira/middleware/service_provider/private.py

@@ -11,7 +11,7 @@
 from config import provider_settings
 from . import resign_provider
 from .schema import BaseProvider, ProviderException
-from ...sdk.endpoint import Driver
+from ...sdk.endpoint.tee import Driver
 
 WHITE_LIST = []
 if provider_settings.get("private", default=None) is not None:
@@ -40,13 +40,13 @@ async def authenticate(self, uid, token, status) -> bool:
         if not Driver.from_public_env().available:
             raise ProviderException(
                 "\nYou are using a public and free instance.\nThe current instance key is not configured.",
-                provider=self.name
+                provider=self.name,
             )
         raise ProviderException(
             "This is a private instance."
             "\nPlease contact the administrator to apply for a private instance."
             f"\n You id is {uid}",
-            provider=self.name
+            provider=self.name,
         )
 
     async def request_driver(self, uid, token) -> Driver:

llmkira/middleware/service_provider/public.py

@@ -12,14 +12,16 @@
 from ...sdk.cache import global_cache_runtime
 from . import resign_provider
 from .schema import BaseProvider, ProviderException
-from ...sdk.endpoint import Driver
+from ...sdk.endpoint.tee import Driver
 
 QUOTA = 24
 WHITE_LIST = []
 if provider_settings.get("public", default=None) is not None:
     QUOTA = provider_settings.public.get("public_quota", default=24)
     WHITE_LIST = provider_settings.public.get("public_white_list", default=[])
-    logger.debug(f"🍦 Public Provider Config Loaded, QUOTA({QUOTA}) WHITE_LIST({WHITE_LIST})")
+    logger.debug(
+        f"🍦 Public Provider Config Loaded, QUOTA({QUOTA}) WHITE_LIST({WHITE_LIST})"
+    )
 
 
 class UserToday(BaseModel):
@@ -42,12 +44,12 @@ async def authenticate(self, uid, token, status) -> bool:
         if not _pass:
             raise ProviderException(
                 "You are using a public instance. You triggered data flood protection today",
-                provider=self.name
+                provider=self.name,
             )
         if not Driver.from_public_env().available:
             raise ProviderException(
                 "You are using a public instance\nBut current instance apikey unavailable",
-                provider=self.name
+                provider=self.name,
             )
         return True
 
@@ -61,14 +63,18 @@ async def check_times(self, times: int, uid: str):
         if read:
             _data: UserToday = UserToday.model_validate(read)
             if str(_data.time) != str(date):
-                await cache.set_data(self.__database_key(uid=uid), value=UserToday().model_dump())
+                await cache.set_data(
+                    self.__database_key(uid=uid), value=UserToday().model_dump()
+                )
                 return True
             else:
                 if _data.count > times:
                     return False
                 if _data.count < times:
                     _data.count += 1
-                    await cache.set_data(self.__database_key(uid=uid), value=_data.model_dump())
+                    await cache.set_data(
+                        self.__database_key(uid=uid), value=_data.model_dump()
+                    )
                     return True
         else:
             _data = UserToday()

llmkira/middleware/service_provider/schema.py

@@ -7,9 +7,10 @@
 
 from pydantic import field_validator, Field
 
-from llmkira.sdk.endpoint import Driver
 from pydantic_settings import BaseSettings
 
+from llmkira.sdk.endpoint.tee import Driver
+
 
 class ProviderSetting(BaseSettings):
     provider: str = Field("PUBLIC", validation_alias="SERVICE_PROVIDER")
@@ -27,7 +28,6 @@ def provider_upper(cls, v):
 
 
 class ProviderException(Exception):
-
     def __init__(self, message: str, provider: str = None):
         self.message = message
         self.provider = provider
@@ -57,7 +57,9 @@ async def authenticate(self, uid, token, status) -> bool:
         """
         必须提供认证文档
         """
-        raise ProviderException("Base Provider auth your token,refer docs", provider=self.name)
+        raise ProviderException(
+            "Base Provider auth your token,refer docs", provider=self.name
+        )
 
     @abstractmethod
     async def request_driver(self, uid, token) -> Driver: