| from typing import Any | |||||
| from typing import Any, Optional | |||||
| import flask_restful | import flask_restful | ||||
| from flask_login import current_user | from flask_login import current_user | ||||
| method_decorators = [account_initialization_required, login_required, setup_required] | method_decorators = [account_initialization_required, login_required, setup_required] | ||||
| resource_type: str | None = None | resource_type: str | None = None | ||||
| resource_model: Any = None | |||||
| resource_model: Optional[Any] = None | |||||
| resource_id_field: str | None = None | resource_id_field: str | None = None | ||||
| token_prefix: str | None = None | token_prefix: str | None = None | ||||
| max_keys = 10 | max_keys = 10 | ||||
| method_decorators = [account_initialization_required, login_required, setup_required] | method_decorators = [account_initialization_required, login_required, setup_required] | ||||
| resource_type: str | None = None | resource_type: str | None = None | ||||
| resource_model: Any = None | |||||
| resource_model: Optional[Any] = None | |||||
| resource_id_field: str | None = None | resource_id_field: str | None = None | ||||
| def delete(self, resource_id, api_key_id): | def delete(self, resource_id, api_key_id): |
| """ | """ | ||||
| event: QueueEvent = QueueEvent.ERROR | event: QueueEvent = QueueEvent.ERROR | ||||
| error: Any = None | |||||
| error: Optional[Any] = None | |||||
| class QueuePingEvent(AppQueueEvent): | class QueuePingEvent(AppQueueEvent): |
| elif isinstance(e, InvokeError | ValueError): | elif isinstance(e, InvokeError | ValueError): | ||||
| err = e | err = e | ||||
| else: | else: | ||||
| err = Exception(e.description if getattr(e, "description", None) is not None else str(e)) | |||||
| description = getattr(e, "description", None) | |||||
| err = Exception(description if description is not None else str(e)) | |||||
| if not message_id or not session: | if not message_id or not session: | ||||
| return err | return err |
| class ModuleExtension(BaseModel): | class ModuleExtension(BaseModel): | ||||
| extension_class: Any = None | |||||
| extension_class: Optional[Any] = None | |||||
| name: str | name: str | ||||
| label: Optional[dict] = None | label: Optional[dict] = None | ||||
| form_schema: Optional[list] = None | form_schema: Optional[list] = None |
| def extension_class(self, module: ExtensionModule, extension_name: str) -> type: | def extension_class(self, module: ExtensionModule, extension_name: str) -> type: | ||||
| module_extension = self.module_extension(module, extension_name) | module_extension = self.module_extension(module, extension_name) | ||||
| assert module_extension.extension_class is not None | |||||
| t: type = module_extension.extension_class | t: type = module_extension.extension_class | ||||
| return t | return t | ||||
| from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError | from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError | ||||
| from datetime import timedelta | from datetime import timedelta | ||||
| from types import TracebackType | from types import TracebackType | ||||
| from typing import Any, Generic, Self, TypeVar | |||||
| from typing import Any, Generic, Optional, Self, TypeVar | |||||
| from httpx import HTTPStatusError | from httpx import HTTPStatusError | ||||
| from pydantic import BaseModel | from pydantic import BaseModel | ||||
| request: SendRequestT, | request: SendRequestT, | ||||
| result_type: type[ReceiveResultT], | result_type: type[ReceiveResultT], | ||||
| request_read_timeout_seconds: timedelta | None = None, | request_read_timeout_seconds: timedelta | None = None, | ||||
| metadata: MessageMetadata = None, | |||||
| metadata: Optional[MessageMetadata] = None, | |||||
| ) -> ReceiveResultT: | ) -> ReceiveResultT: | ||||
| """ | """ | ||||
| Sends a request and wait for a response. Raises an McpError if the | Sends a request and wait for a response. Raises an McpError if the |
| """A message with specific metadata for transport-specific features.""" | """A message with specific metadata for transport-specific features.""" | ||||
| message: JSONRPCMessage | message: JSONRPCMessage | ||||
| metadata: MessageMetadata = None | |||||
| metadata: Optional[MessageMetadata] = None | |||||
| class OAuthClientMetadata(BaseModel): | class OAuthClientMetadata(BaseModel): |
| import logging | import logging | ||||
| from threading import Lock | from threading import Lock | ||||
| from typing import Any | |||||
| from typing import Any, Optional | |||||
| logger = logging.getLogger(__name__) | logger = logging.getLogger(__name__) | ||||
| _tokenizer: Any = None | |||||
| _tokenizer: Optional[Any] = None | |||||
| _lock = Lock() | _lock = Lock() | ||||
| from collections.abc import Generator | from collections.abc import Generator | ||||
| from datetime import datetime | from datetime import datetime | ||||
| from typing import Any | |||||
| from typing import Any, Optional | |||||
| from core.rag.extractor.watercrawl.client import WaterCrawlAPIClient | from core.rag.extractor.watercrawl.client import WaterCrawlAPIClient | ||||
| def __init__(self, api_key, base_url: str | None = None): | def __init__(self, api_key, base_url: str | None = None): | ||||
| self.client = WaterCrawlAPIClient(api_key, base_url) | self.client = WaterCrawlAPIClient(api_key, base_url) | ||||
| def crawl_url(self, url, options: dict | Any = None) -> dict: | |||||
| def crawl_url(self, url, options: Optional[dict | Any] = None) -> dict: | |||||
| options = options or {} | options = options or {} | ||||
| spider_options = { | spider_options = { | ||||
| "max_depth": 1, | "max_depth": 1, |
| import ssl | import ssl | ||||
| from collections.abc import Callable | from collections.abc import Callable | ||||
| from datetime import timedelta | from datetime import timedelta | ||||
| from typing import TYPE_CHECKING, Any, Union | |||||
| from typing import TYPE_CHECKING, Any, Optional, Union | |||||
| import redis | import redis | ||||
| from redis import RedisError | from redis import RedisError | ||||
| app.extensions["redis"] = redis_client | app.extensions["redis"] = redis_client | ||||
| def redis_fallback(default_return: Any = None): | |||||
| def redis_fallback(default_return: Optional[Any] = None): | |||||
| """ | """ | ||||
| decorator to handle Redis operation exceptions and return a default value when Redis is unavailable. | decorator to handle Redis operation exceptions and return a default value when Redis is unavailable. | ||||
| import logging | import logging | ||||
| from typing import Optional | |||||
| from core.tools.entities.api_entities import ToolProviderTypeApiLiteral | from core.tools.entities.api_entities import ToolProviderTypeApiLiteral | ||||
| from core.tools.tool_manager import ToolManager | from core.tools.tool_manager import ToolManager | ||||
| class ToolCommonService: | class ToolCommonService: | ||||
| @staticmethod | @staticmethod | ||||
| def list_tool_providers(user_id: str, tenant_id: str, typ: ToolProviderTypeApiLiteral = None): | |||||
| def list_tool_providers(user_id: str, tenant_id: str, typ: Optional[ToolProviderTypeApiLiteral] = None): | |||||
| """ | """ | ||||
| list tool providers | list tool providers | ||||
| ) | ) | ||||
| role_prefix = None | role_prefix = None | ||||
| prompts: Any = None | |||||
| prompts: Optional[Any] = None | |||||
| # Chat Model | # Chat Model | ||||
| if model_config.mode == LLMMode.CHAT.value: | if model_config.mode == LLMMode.CHAT.value: |
| import os | import os | ||||
| from collections import UserDict | from collections import UserDict | ||||
| from typing import Optional | |||||
| from unittest.mock import MagicMock | from unittest.mock import MagicMock | ||||
| import pytest | import pytest | ||||
| def mock_vector_db_client( | def mock_vector_db_client( | ||||
| self, | self, | ||||
| config=None, | config=None, | ||||
| adapter: HTTPAdapter = None, | |||||
| adapter: Optional[HTTPAdapter] = None, | |||||
| ): | ): | ||||
| self.conn = MagicMock() | self.conn = MagicMock() | ||||
| self._config = MagicMock() | self._config = MagicMock() |
| key="", | key="", | ||||
| read_consistency: ReadConsistency = ReadConsistency.EVENTUAL_CONSISTENCY, | read_consistency: ReadConsistency = ReadConsistency.EVENTUAL_CONSISTENCY, | ||||
| timeout=10, | timeout=10, | ||||
| adapter: HTTPAdapter = None, | |||||
| adapter: Optional[HTTPAdapter] = None, | |||||
| pool_size: int = 2, | pool_size: int = 2, | ||||
| proxies: Optional[dict] = None, | proxies: Optional[dict] = None, | ||||
| password: Optional[str] = None, | password: Optional[str] = None, | ||||
| shard: int, | shard: int, | ||||
| replicas: int, | replicas: int, | ||||
| description: Optional[str] = None, | description: Optional[str] = None, | ||||
| index: Index = None, | |||||
| embedding: Embedding = None, | |||||
| index: Optional[Index] = None, | |||||
| embedding: Optional[Embedding] = None, | |||||
| timeout: Optional[float] = None, | timeout: Optional[float] = None, | ||||
| ttl_config: Optional[dict] = None, | ttl_config: Optional[dict] = None, | ||||
| filter_index_config: FilterIndexConfig = None, | |||||
| filter_index_config: Optional[FilterIndexConfig] = None, | |||||
| indexes: Optional[list[IndexField]] = None, | indexes: Optional[list[IndexField]] = None, | ||||
| ) -> RPCCollection: | ) -> RPCCollection: | ||||
| return RPCCollection( | return RPCCollection( | ||||
| database_name: str, | database_name: str, | ||||
| collection_name: str, | collection_name: str, | ||||
| vectors: list[list[float]], | vectors: list[list[float]], | ||||
| filter: Filter = None, | |||||
| filter: Optional[Filter] = None, | |||||
| params=None, | params=None, | ||||
| retrieve_vector: bool = False, | retrieve_vector: bool = False, | ||||
| limit: int = 10, | limit: int = 10, | ||||
| collection_name: str, | collection_name: str, | ||||
| ann: Optional[Union[list[AnnSearch], AnnSearch]] = None, | ann: Optional[Union[list[AnnSearch], AnnSearch]] = None, | ||||
| match: Optional[Union[list[KeywordSearch], KeywordSearch]] = None, | match: Optional[Union[list[KeywordSearch], KeywordSearch]] = None, | ||||
| filter: Union[Filter, str] = None, | |||||
| filter: Optional[Union[Filter, str]] = None, | |||||
| rerank: Optional[Rerank] = None, | rerank: Optional[Rerank] = None, | ||||
| retrieve_vector: Optional[bool] = None, | retrieve_vector: Optional[bool] = None, | ||||
| output_fields: Optional[list[str]] = None, | output_fields: Optional[list[str]] = None, | ||||
| database_name: str, | database_name: str, | ||||
| collection_name: str, | collection_name: str, | ||||
| document_ids: Optional[list[str]] = None, | document_ids: Optional[list[str]] = None, | ||||
| filter: Filter = None, | |||||
| filter: Optional[Filter] = None, | |||||
| timeout: Optional[float] = None, | timeout: Optional[float] = None, | ||||
| ): | ): | ||||
| return {"code": 0, "msg": "operation success"} | return {"code": 0, "msg": "operation success"} |