fix MCP asyncio error

Author: nicucalcea

augmenta/agent.py

@@ -119,6 +119,14 @@ def create_structure_class(yaml_file_path: Union[str, Path]) -> Type[BaseModel]:
 
         except (yaml.YAMLError, OSError) as e:
             raise ValueError(f"Failed to parse YAML: {e}")
+    
+    def get_mcp_servers_context(self):
+        """Get the MCP servers context manager from the underlying agent.
+        
+        Returns:
+            Context manager for running MCP servers
+        """
+        return self.agent.run_mcp_servers()
 
     async def run(
         self,
@@ -147,12 +155,12 @@ async def run(
             if temperature is not None and temperature != self.temperature:
                 model_settings = self._create_model_settings(temperature)
 
-            async with self.agent.run_mcp_servers():
-                result = await self.agent.run(
-                    prompt,
-                    output_type=response_format,
-                    model_settings=model_settings
-                )
+            # Note: MCP servers context should be managed at a higher level now
+            result = await self.agent.run(
+                prompt,
+                output_type=response_format,
+                model_settings=model_settings
+            )
 
             # Return the appropriate result format
             return result.data.model_dump() if response_format else result.data

augmenta/augmenta.py

@@ -165,8 +165,10 @@ async def process_with_limit(row: Dict[str, Any]) -> ProcessingResult:
                 progress_callback=update_progress
             )
 
-    tasks = [process_with_limit(row) for row in rows_to_process]
-    results = await asyncio.gather(*tasks)
+    # Run all tasks under a single MCP server context
+    async with agent.get_mcp_servers_context():
+        tasks = [process_with_limit(row) for row in rows_to_process]
+        results = await asyncio.gather(*tasks)
 
     # Update DataFrame with results
     successful_results, error_count = update_dataframe_with_results(df, results)

docs/examples/donations/README.md

@@ -62,7 +62,6 @@ output_csv: data/donations_classified.csv
 model:
   name: openai/gpt-4o-mini
   max_tokens: 20000
-query_col: DonorName
 search:
   engine: brave
   country: GB

docs/examples/donations/README.qmd

@@ -55,7 +55,6 @@ output_csv: data/donations_classified.csv
 model:
   name: openai/gpt-4o-mini
   max_tokens: 20000
-query_col: DonorName
 search:
   engine: brave
   country: GB

docs/examples/donations/config.yaml

@@ -4,7 +4,6 @@ agent: autonomous
 model:
   provider: openai
   name: gpt-4o-mini
-query_col: DonorName
 search:
   engine: duckduckgo
   results: 20

docs/examples/maps/config.yaml

@@ -4,7 +4,6 @@ agent: autonomous
 model:
   provider: openai
   name: gpt-4o-mini
-query_col: city
 search:
   engine: duckduckgo
   results: 20

docs/search.md

@@ -1,6 +1,6 @@
 # Searching with Augmenta
 
-Augmenta supports several search engines out of the box. It will search for the `query_col` you specify, extract the links, and return the text from those links.
+Augmenta supports several search engines out of the box, with more available as [MCP servers](/docs/tools.md).
 
 Some search engines have a rate limit, meaning they will only allow a certain number of searches in a given time period. You can configure Augmenta to respect this rate limit by setting the `rate_limit` parameter in your configuration file.
 

docs/tools.md

@@ -1,6 +1,6 @@
 # Tools
 
-By default, Augmenta has access to two tools: a search engine and a tool for retrieving text from the internet.
+By default, Augmenta has access to three tools: a search engine, a tool for retrieving text from the internet, and a tool for attaching local files.
 
 However, you may be interested in empowering it to do more. For example, you may want to give it access to a database, or give it the power to write code or do complex calculations.
 
@@ -34,4 +34,5 @@ Note that the Google Maps server needs a `GOOGLE_MAPS_API_KEY` environment varia
 GOOGLE_MAPS_API_KEY=XXXXX
 ```
 
-Augmenta will now have access to Google Maps. However, it will still have access to the default search tool, meaning it may decide to use it instead of the Google Maps server. You can explicitly tell the LLM to use the newly added tool in the prompts.
+Augmenta will now have access to Google Maps. However, it will still have access to the default search tool, meaning it may decide to use it instead of the Google Maps server. You can explicitly tell the LLM to use the newly added tool in the prompts.
+