Merge pull request #375 from LlmKira/dev

Fix import

README.md

@@ -32,7 +32,7 @@
   &
   <a href="README_EN.md">📝 English Readme</a>
   &
-  <a href="CONTRIBUTING.md">🤝 贡献必看</a>
+  <a href=".github/CONTRIBUTING.md">🤝 贡献必看</a>
 </p>
 
 

llmkira/extra/user/schema.py

@@ -11,7 +11,7 @@
 from pydantic import field_validator, ConfigDict, BaseModel, Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
-from ...sdk.endpoint import Driver
+from llmkira.sdk.endpoint.tee import Driver
 
 
 class UserDriverMode(Enum):
@@ -74,7 +74,7 @@ def create_from_function(
         return cls(
             request_id=request_id,
             uid=uid,
-            cost=cls.Cost.by_function(
+            cost=Cost.by_function(
                 function_name=cost_by,
                 token_usage=token_usage,
                 token_uuid=token_uuid,

llmkira/middleware/llm_provider.py

@@ -5,9 +5,10 @@
 # @Software: PyCharm
 from loguru import logger
 
-from llmkira.extra.user import UserControl, UserConfig, UserDriverMode
+from llmkira.extra.user import UserControl, UserConfig
 from llmkira.middleware.service_provider.schema import ProviderSettingObj
 from .service_provider import loaded_provider, PublicProvider
+from ..extra.user.schema import UserDriverMode
 
 if not loaded_provider:
     raise Exception("⚠️ No Any Driver Provider Loaded, Even Public Provider")
@@ -34,27 +35,33 @@ async def get(self):
         :return: Driver
         """
         self.user = await UserControl.get_driver_config(uid=self.uid)
-        assert isinstance(self.user, UserConfig.LlmConfig), "UserConfig.LlmConfig is empty"
+        assert isinstance(
+            self.user, UserConfig.LlmConfig
+        ), "UserConfig.LlmConfig is empty"
         # 配置了自己的私有例
         if self.user.mode == UserDriverMode.private:
             return self.user.driver
         # Public Provider
         if ProviderSettingObj.is_open_everyone:
             provider = PublicProvider()
-            logger.debug(f"🍦 Public Provider --name ({provider.name}) --mode ({self.user.mode}) --uid ({self.uid})")
+            logger.debug(
+                f"🍦 Public Provider --name ({provider.name}) --mode ({self.user.mode}) --uid ({self.uid})"
+            )
             if await provider.authenticate(
-                    uid=self.uid,
-                    token=self.user.token, status=self.user.mode):
-                return await provider.request_driver(uid=self.uid, token=self.user.token)
+                uid=self.uid, token=self.user.token, status=self.user.mode
+            ):
+                return await provider.request_driver(
+                    uid=self.uid, token=self.user.token
+                )
         else:
             # 用户需要特别配置 Token
             provider = loaded_provider()
             if await provider.authenticate(
-                    uid=self.uid,
-                    token=self.user.token,
-                    status=self.user.mode
+                uid=self.uid, token=self.user.token, status=self.user.mode
             ):
-                return await provider.request_driver(uid=self.uid, token=self.user.token)
+                return await provider.request_driver(
+                    uid=self.uid, token=self.user.token
+                )
         """
         raise ProviderException(
             f"AuthChanged {self.user.provider} >>change>> {loaded_provider.name.upper()}"

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()