kgc-agents / agents.py
RubenPeeters
Change agents
47a0a0e
import os
from smolagents import CodeAgent, OpenAIServerModel
from tools import (
validate_rdf_syntax,
get_entities_from_kg,
get_types_from_ontology,
get_relations_from_ontology,
)
from dotenv import load_dotenv
load_dotenv()
# model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
# model = TransformersModel(model_id=model_id)
model = OpenAIServerModel(
model_id="qwen/qwen-2.5-coder-32b-instruct:free",
# model_id="microsoft/mai-ds-r1:free",
api_base="https://openrouter.ai/api/v1",
api_key=os.environ["OPENROUTER_API_KEY"],
)
# web_agent = CodeAgent(
# tools=[DuckDuckGoSearchTool()],
# model=model,
# name="web_search_agent",
# description="Has the tools to run web searches for you. Give it your query as an argument."
# # prompt_templates=PromptTemplates(system_prompt="")
# )
kgc_agent = CodeAgent(
tools=[get_types_from_ontology, get_entities_from_kg, get_relations_from_ontology],
additional_authorized_imports=["rdflib"],
model=model,
name="named_entity_recognition_expert",
description="An agent that is an expert in Named Entity Recognition. Named Entity Recognition (NER) is a subfield of computer science and Natural Language Processing (NLP) that focuses on identifying and classifying entities in unstructured text into predefined categories, such as persons, geographical locations and organizations. Provide content to this agent to detect the entities inside.",
)
validation_agent = CodeAgent(
tools=[validate_rdf_syntax],
additional_authorized_imports=["rdflib"],
model=model,
name="knowledge_graph_expert",
description="Has extensive knowledge on knowledge graphs. Has the final say before writing the output to a file.",
)
# TODO: Add RDF specification tool, CACAO ontology tool, reasoning tool?
ner_agent = CodeAgent(
tools=[],
model=model,
name="named_entity_recognition_expert",
description="An agent that is an expert in Named Entity Recognition. Named Entity Recognition (NER) is a subfield of computer science and Natural Language Processing (NLP) that focuses on identifying and classifying entities in unstructured text into predefined categories, such as persons, geographical locations and organizations. Provide content to this agent to detect the entities inside.",
)
et_agent = CodeAgent(
tools=[get_types_from_ontology],
additional_authorized_imports=["rdflib"],
model=model,
name="entity_typing_expert",
description="An agent that is an expert in Entity Typing. tasks provide fine-grained and ultra-grained type information for entities such as scientists, clubs, and hotels and is the process of categorizing named entities into predefined types. Provide named entities to this agent and expect typed entities back.",
)
el_agent = CodeAgent(
tools=[get_entities_from_kg],
model=model,
additional_authorized_imports=["rdflib"],
name="entity_linking_expert",
description="An agent that is an expert in Entity Linking. Entity Linking (EL) tasks, or entity disambiguation, involves identifying named entities within a text and linking them to their corresponding entries in a knowledge graph. Provide typed entities to this agent to link them to an existing knowledge base.",
)
co_agent = CodeAgent(
tools=[],
additional_authorized_imports=["rdflib"],
model=model,
name="coreference_resolution_expert",
description="An agent that is an expert in Coreference Resolution. Coreference Resolution (CO) focuses on determining when two or more expressions refer to the same entity. Provide texts to this agent to find out which texts refer to the same entities.",
)
re_agent = CodeAgent(
tools=[get_relations_from_ontology],
model=model,
additional_authorized_imports=["rdflib"],
name="relation_extraction_expert",
description="An agent that is an expert in Relation Extraction. Relation Extraction (RE) involves identifying facts about relations between extracted entities. Provide a text and entities to this agent to find relations between the entities from the text.",
)
# ch_agent = CodeAgent(
# tools=[DuckDuckGoSearchTool()],
# additional_authorized_imports=["rdflib"],
# model=model,
# name="cultural_heritage_domain_expert",
# description="Has extensive knowledge of cultural heritage. Ask it about the correctness of facts and the interpretation of a knowledge graph. It also has access to the internet to check facts.",
# )
kg_agent = CodeAgent(
tools=[validate_rdf_syntax],
additional_authorized_imports=["rdflib"],
model=model,
name="knowledge_graph_expert",
description="Has extensive knowledge on knowledge graphs. Has the final say before writing the output to a file.",
)
manager_agent = CodeAgent(
# Using N = 2, following the example of the HF team: https://huggingface.co/blog/beating-gaia
tools=[],
additional_authorized_imports=["rdflib"],
model=model,
managed_agents=[ner_agent, et_agent, el_agent, co_agent, re_agent],
planning_interval=2,
)