Merge pull request #124 from UiPath/fix/linter

fix: lint types

Author: cristipufu

pyproject.toml

@@ -49,6 +49,7 @@ dev = [
     "pytest>=7.4.0",
     "pytest-cov>=4.1.0",
     "pytest-mock>=3.11.1",
+    "pytest-asyncio>=1.0.0",
     "pre-commit>=4.1.0",
     "numpy>=1.24.0",
     "pytest_httpx>=0.35.0"
@@ -100,6 +101,8 @@ disallow_untyped_defs = false
 testpaths = ["tests"]
 python_files = "test_*.py"
 addopts = "-ra -q"
+asyncio_default_fixture_loop_scope = "function"
+asyncio_mode = "auto"
 
 [[tool.uv.index]]
 name = "testpypi"

src/uipath_langchain/_cli/_runtime/_context.py

@@ -11,7 +11,7 @@ class LangGraphRuntimeContext(UiPathRuntimeContext):
     """Context information passed throughout the runtime execution."""
 
     langgraph_config: Optional[LangGraphConfig] = None
-    state_graph: Optional[StateGraph] = None
+    state_graph: Optional[StateGraph[Any, Any]] = None
     output: Optional[Any] = None
     state: Optional[Any] = (
         None  # TypedDict issue, the actual type is: Optional[langgraph.types.StateSnapshot]

src/uipath_langchain/_cli/_runtime/_runtime.py

@@ -271,7 +271,9 @@ async def cleanup(self):
         if hasattr(self, "graph_config") and self.graph_config:
             await self.graph_config.cleanup()
 
-    def _extract_graph_result(self, final_chunk, graph: CompiledStateGraph):
+    def _extract_graph_result(
+        self, final_chunk, graph: CompiledStateGraph[Any, Any, Any]
+    ):
         """
         Extract the result from a LangGraph output chunk according to the graph's output channels.
 

src/uipath_langchain/_cli/_utils/_graph.py

@@ -21,14 +21,18 @@ class GraphConfig:
     path: str
     file_path: str
     graph_var: str
-    _graph: Optional[Union[StateGraph, CompiledStateGraph]] = None
+    _graph: Optional[Union[StateGraph[Any, Any], CompiledStateGraph[Any, Any, Any]]] = (
+        None
+    )
 
     @classmethod
     def from_config(cls, name: str, path: str) -> "GraphConfig":
         file_path, graph_var = path.split(":")
         return cls(name=name, path=path, file_path=file_path, graph_var=graph_var)
 
-    async def load_graph(self) -> Union[StateGraph, CompiledStateGraph]:
+    async def load_graph(
+        self,
+    ) -> Union[StateGraph[Any, Any], CompiledStateGraph[Any, Any, Any]]:
         """Load graph from the specified path"""
         try:
             cwd = os.path.abspath(os.getcwd())

src/uipath_langchain/_cli/cli_init.py

@@ -54,7 +54,9 @@ def process_nullable_types(
     return schema
 
 
-def generate_schema_from_graph(graph: CompiledStateGraph) -> Dict[str, Any]:
+def generate_schema_from_graph(
+    graph: CompiledStateGraph[Any, Any, Any],
+) -> Dict[str, Any]:
     """Extract input/output schema from a LangGraph graph"""
     schema = {
         "input": {"type": "object", "properties": {}, "required": []},

src/uipath_langchain/embeddings/embeddings.py

@@ -59,7 +59,7 @@ class UiPathOpenAIEmbeddings(UiPathRequestMixin, OpenAIEmbeddings):
     )
 
     def embed_documents(
-        self, texts: List[str], chunk_size: Optional[int] = None
+        self, texts: List[str], chunk_size: Optional[int] = None, **kwargs
     ) -> List[List[float]]:
         """Embed a list of documents using the UiPath."""
         embeddings = []
@@ -83,6 +83,7 @@ async def aembed_documents(
         self,
         texts: List[str],
         chunk_size: Optional[int] = None,
+        **kwargs,
     ) -> List[List[float]]:
         """Async version of embed_documents."""
         embeddings = []

tests/hitl/mocks/api_trigger_hitl.py

@@ -1,26 +1,20 @@
-import dataclasses
+from typing import TypedDict
 
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import END, START, StateGraph
 from langgraph.types import interrupt
 
 
-@dataclasses.dataclass
-class Input:
-    pass
-
-
-@dataclasses.dataclass
-class Output:
+class State(TypedDict):
     message: str
 
 
-def main_node(input: Input) -> Output:
+def main_node(state: State) -> State:
     response = interrupt("interrupt message")
-    return Output(message=response)
+    return {"message": response}
 
 
-builder = StateGraph(input=Input, output=Output)
+builder: StateGraph[State] = StateGraph(State)
 
 builder.add_node("main_node", main_node)