diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/__init__.py b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/__init__.py index 28267eec885d..3d73e502f611 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/__init__.py @@ -1,5 +1,4 @@ from ._config import ( - EmbeddingConfig, GlobalContextConfig, GlobalDataConfig, LocalContextConfig, @@ -23,5 +22,4 @@ "LocalSearchToolReturn", "MapReduceConfig", "SearchConfig", - "EmbeddingConfig", ] diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_config.py b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_config.py index d6adf65e39c5..10c3d4a985e6 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_config.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_config.py @@ -1,5 +1,3 @@ -from typing import Callable, Literal, Optional - from pydantic import BaseModel @@ -59,16 +57,3 @@ class SearchConfig(BaseModel): max_tokens: int = 1500 temperature: float = 0.0 response_type: str = "multiple paragraphs" - - -class EmbeddingConfig(BaseModel): - api_key: Optional[str] = None - model: str - api_base: Optional[str] = None - deployment_name: Optional[str] = None - api_version: Optional[str] = None - api_type: Literal["azure", "openai"] = "openai" - organization: Optional[str] = None - azure_ad_token_provider: Optional[Callable[[], str]] = None - max_retries: int = 10 - request_timeout: float = 180.0 diff --git a/python/packages/autogen-ext/tests/tools/test_graphrag_tools.py b/python/packages/autogen-ext/tests/tools/test_graphrag_tools.py index 324c0e1b6741..1c6f662a2834 100644 --- a/python/packages/autogen-ext/tests/tools/test_graphrag_tools.py +++ b/python/packages/autogen-ext/tests/tools/test_graphrag_tools.py @@ -39,7 +39,7 @@ async def agenerate( ) -> str: return "Mock response" - async def astream_generate( + async def astream_generate( # type: ignore self, messages: str | list[Any], callbacks: list[BaseLLMCallback] | None = None, **kwargs: Any ) -> AsyncGenerator[str, None]: yield "Mock response"