mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 23:47:46 +00:00
you-n-g
4f0b2be7a7
* feat: init pydantic ai agent & context 7 mcp * feat: integrate MCP documentation search into data science pipeline evaluation * fix: disable MCP documentation search and update related docstrings and defaults * lint * fix: correct prompt formatting and conditional blocks in pipeline_eval section * lint * feat: add query method to PAIAgent for synchronous agent execution * fix: apply nest_asyncio for agent and update context7 query method * lint * lint * lint * lint * docs: update MCP folder docstring and rename test class in test_pydantic.py * refactor: centralize completion kwargs logic and update pydantic_ai integration * fixbug * typo * fix: bug triggered by padantic-ai version backtracking. --------- Co-authored-by: Linlang <Lv.Linlang@hotmail.com>
45 lines
1.1 KiB
Python
45 lines
1.1 KiB
Python
from abc import abstractmethod
|
|
|
|
import nest_asyncio
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.mcp import MCPServerStreamableHTTP
|
|
|
|
from rdagent.oai.backend.pydantic_ai import get_agent_model
|
|
|
|
|
|
class BaseAgent:
|
|
|
|
@abstractmethod
|
|
def __init__(self, system_prompt: str, toolsets: list[str]): ...
|
|
|
|
@abstractmethod
|
|
def query(self, query: str) -> str: ...
|
|
|
|
|
|
class PAIAgent(BaseAgent):
|
|
"""
|
|
Pydantic-AI agent
|
|
"""
|
|
|
|
agent: Agent
|
|
|
|
def __init__(self, system_prompt: str, toolsets: list[str | MCPServerStreamableHTTP]):
|
|
toolsets = [(ts if isinstance(ts, MCPServerStreamableHTTP) else MCPServerStreamableHTTP(ts)) for ts in toolsets]
|
|
self.agent = Agent(get_agent_model(), system_prompt=system_prompt, toolsets=toolsets)
|
|
|
|
def query(self, query: str) -> str:
|
|
"""
|
|
|
|
Parameters
|
|
----------
|
|
query : str
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
"""
|
|
|
|
nest_asyncio.apply() # NOTE: very important. Because pydantic-ai uses asyncio!
|
|
result = self.agent.run_sync(query)
|
|
return result.output
|