| 8 | client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) |
| 9 | |
| 10 | def extract_nodes(data: any, current_nodes: any) -> str: |
| 11 | try: |
| 12 | response = client.chat.completions.create( |
| 13 | messages=[ |
| 14 | { |
| 15 | "role": "system", |
| 16 | "content": """ |
| 17 | You are an AI assistant that extracts new entity nodes from text. Your task is to identify and extract significant entities mentioned in the tweet, focusing on people, ideas, and concepts. |
| 18 | |
| 19 | Return the extracted nodes in the following JSON format: |
| 20 | { |
| 21 | "nodes": [ |
| 22 | { "label": "string" } |
| 23 | ] |
| 24 | } |
| 25 | |
| 26 | Ensure that each node has a concise label that provides more context about the node. |
| 27 | |
| 28 | Rules: |
| 29 | - Only extract nodes that are relevant to the tweet. |
| 30 | - Only extract nodes that are not already in the current nodes. |
| 31 | - Only extract 5 nodes MAXIUMUM |
| 32 | """ |
| 33 | }, |
| 34 | { |
| 35 | "role": "user", |
| 36 | "content": f""" |
| 37 | Given the current nodes, extract new nodes from the following text: |
| 38 | |
| 39 | Current nodes: |
| 40 | |
| 41 | {current_nodes} |
| 42 | |
| 43 | Current Text: |
| 44 | |
| 45 | {data['text']} |
| 46 | |
| 47 | Consider people, ideas, and concepts. Avoid creating nodes for actions or temporal information. |
| 48 | """ |
| 49 | } |
| 50 | ], |
| 51 | model="gpt-4o-mini", |
| 52 | response_format={ "type": "json_object" } |
| 53 | ) |
| 54 | |
| 55 | nodes = response.choices[0].message.content |
| 56 | nodes = json.loads(nodes)['nodes'] |
| 57 | |
| 58 | nodes = [{"id": f"{data['id']}_{i}", "text": data['text'], **node} for i, node in enumerate(nodes)] |
| 59 | |
| 60 | return nodes |
| 61 | except Exception as e: |
| 62 | logging.error(f"Error extracting nodes: {str(e)}") |
| 63 | return [] |
| 64 | |
| 65 | def extract_edges(data: dict) -> list: |
| 66 | try: |