import mcp
from langchain.agents import initialize_agent, Tool
from langchain_openai import ChatOpenAI
from langchain.tools import FileManagementTool
from datetime import datetime
# Configuration
MCP_ENDPOINT = "<<MCP_ENDPOINT>>"
CORVIC_API_TOKEN = "<<YOUR_CORVIC_API_TOKEN>>"
OPENAI_API_KEY = "<<YOUR_OPENAI_API_KEY>>"
headers = {
"Authorization": f"Bearer {CORVIC_API_TOKEN}",
"Content-Type": "application/json"
}
# Create Corvic query function
async def query_corvic(query: str) -> str:
"""Query Corvic agent using MCP."""
async with mcp.ClientSession(
transport=mcp.SSEClientTransport(MCP_ENDPOINT, headers=headers)
) as session:
await session.initialize()
query_result = await session.call_tool(
"query",
arguments={"query_content": query}
)
return query_result.content[0].text if query_result.content else ""
# Create tools
corvic_tool = Tool(
name="Corvic Query",
func=lambda q: asyncio.run(query_corvic(q)),
description="Query Corvic agent for NAICS code information"
)
file_tool = FileManagementTool()
# Initialize LangChain agent
llm = ChatOpenAI(model="gpt-4", temperature=0, api_key=OPENAI_API_KEY)
agent = initialize_agent(
tools=[corvic_tool, file_tool],
llm=llm,
agent="zero-shot-react-description",
verbose=True
)
# Execute query
question = """As per NAICS codes, Describe wheat farming. Use tools to answer this question.
The answer should be written to a file named results.md
Write the date and time of the response in the file."""
result = agent.run(question)
# Add timestamp to file
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
with open("results.md", "a") as f:
f.write(f"\n\n---\nResponse generated at: {timestamp}\n")
print(f"Result saved to results.md")
