+ linting and type fixes

microsoft · Jan 9, 2025 · dff93ea · dff93ea
1 parent 2ac3edf
commit dff93ea
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 18 deletions.
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
@@ -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
@@ -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
@@ -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"