wehub-resource-sync
e768098d0e
tools_continuous_delivery / Private PyPI non-main branch release (push) Has been skipped
tools_continuous_delivery / Private PyPI main branch release (push) Failing after 2m42s
Publish Promptflow Doc / Build (push) Has been cancelled
Publish Promptflow Doc / Deploy (push) Has been cancelled
Flake8 Lint / flake8 (push) Has been cancelled
Spell check CI / Spell_Check (push) Has been cancelled
91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
import asyncio
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import List
|
|
|
|
from dotenv import load_dotenv
|
|
from typing_extensions import Never
|
|
|
|
from agent_framework import Agent, Executor, WorkflowBuilder, WorkflowContext, handler
|
|
from agent_framework.openai import OpenAIChatClient
|
|
|
|
load_dotenv()
|
|
|
|
_NER_SYSTEM_PROMPT = (
|
|
"Your task is to find entities of certain type from the given text content.\n"
|
|
"If there're multiple entities, please return them all with comma separated, "
|
|
'e.g. "entity1, entity2, entity3".\n'
|
|
"You should only return the entity list, nothing else.\n"
|
|
'If there\'s no such entity, please return "None".'
|
|
)
|
|
|
|
_NER_USER_TEMPLATE = "Entity type: {entity_type}\nText content: {text}\nEntities:"
|
|
|
|
|
|
@dataclass
|
|
class NERInput:
|
|
text: str
|
|
entity_type: str = "job title"
|
|
|
|
|
|
class NERExecutor(Executor):
|
|
def __init__(self, **kwargs):
|
|
super().__init__(**kwargs)
|
|
client = OpenAIChatClient(
|
|
azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
|
|
model=os.environ["AZURE_OPENAI_DEPLOYMENT"],
|
|
api_key=os.environ["AZURE_OPENAI_API_KEY"],
|
|
)
|
|
self._agent = Agent(
|
|
client=client,
|
|
name="NERAgent",
|
|
instructions=_NER_SYSTEM_PROMPT,
|
|
)
|
|
|
|
@handler
|
|
async def extract(self, ner_input: NERInput, ctx: WorkflowContext[str]) -> None:
|
|
user_msg = _NER_USER_TEMPLATE.format(
|
|
entity_type=ner_input.entity_type, text=ner_input.text
|
|
)
|
|
response = await self._agent.run(user_msg)
|
|
await ctx.send_message(response.text)
|
|
|
|
|
|
class CleansingExecutor(Executor):
|
|
@handler
|
|
async def cleanse(self, entities_str: str, ctx: WorkflowContext[Never, List[str]]) -> None:
|
|
parts = entities_str.split(",")
|
|
cleaned = [p.strip(' \t."') for p in parts]
|
|
entities = [p for p in cleaned if len(p) > 0]
|
|
await ctx.yield_output(entities)
|
|
|
|
|
|
def create_workflow():
|
|
"""Create a fresh workflow instance.
|
|
|
|
MAF workflows do not support concurrent execution, so each
|
|
concurrent caller needs its own workflow instance.
|
|
"""
|
|
_ner = NERExecutor(id="NER_LLM")
|
|
_cleansing = CleansingExecutor(id="cleansing")
|
|
return (
|
|
WorkflowBuilder(name="NERWorkflow", start_executor=_ner)
|
|
.add_edge(_ner, _cleansing)
|
|
.build()
|
|
)
|
|
|
|
|
|
async def main():
|
|
workflow = create_workflow()
|
|
result = await workflow.run(
|
|
NERInput(
|
|
text="Maxime is a data scientist at Auto Dataset and he lives in Paris, France.",
|
|
entity_type="job title",
|
|
)
|
|
)
|
|
print(result.get_outputs()[0])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|