🚗 Visualising partnerships on the self-driving car market with an AI agent

AI has the capacity to transform text into structured data to create data visualisations.
Our objective is to create a network graph of partnerships between companies in the autonomous vehicle market using an AI agent. This agent will autonomously identify collaborations from news articles and manage the graph database.
Methodology:
Collecting news articles related to partnerships on the self-driving car market
Setting up an extractor to obtain structured data on partnerships
Creating an agent and tools
Executing the AI agent on articles and monitoring the results
import pandas as pd
import json
import os

from pydantic import BaseModel, Field
from typing import Optional, List
from enum import Enum
from dotenv import load_dotenv

from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
from langchain.agents import create_tool_calling_agent, AgentExecutor

# The dataset
news = pd.read_csv("news.csv")

#The extractor
class Category(str, Enum):
   self_driving_startup = "self driving startup"
   automaker = "automaker"
   software_hardware = "software/hardware"
   transport_delivery = "transport/delivery"
   automotive_supplier = "automotive supplier"
   other = "other"

class Company(BaseModel):
  """Information about a company on the self-driving car market"""

  name: Optional[str] = Field(default=None, description="The name of the company")
  category: Optional[Category] = Field(default=None, description="The type of company on the self-driving car market")


class Partnership(BaseModel):
    """Information about a partnership"""

    companies: List[Company] = Field(default_factory=list, min_items=2, description="The partnership stakeholders")

class Data(BaseModel):
   """Extracted data about partnerships"""

   partnerships: List[Partnership] = Field(default_factory=list)

extract_prompt_template = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are an expert extraction algorithm. "
            "Only extract partnerships between companies on the self-driving car market from the text. "
            "If you do not know the value of an attribute asked to extract, "
            "return null for the attribute's value.",
        ),
        ("human", "{text}"),
    ]
)

load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
llm = init_chat_model("gpt-4o", model_provider="openai")
structured_llm = llm.with_structured_output(schema=Data)

# Function to load the data from JSON file
def load_json() -> dict:
  with open("data.json", "r") as data_json:
    data = json.load(data_json)

  return data

# Function to save the data into the JSON file
def save_json(dict: dict):
  dict_json = json.dumps(dict)
  with open("data.json", "w") as data_json:
    data_json.write(dict_json)


# The tools
@tool
def extract_partnerships(text: str) -> tuple[str, list]:
    """Extract partnerships from self-driving car market text."""

    extract_prompt = extract_prompt_template.invoke({"text": text})
    result = structured_llm.invoke(extract_prompt)

    partnerships = []

    for partnership in result.partnerships:
      partnership_data = {
            "companies": [
                {
                    "name": company.name,
                    "category": company.category.value if company.category else None
                }
                for company in partnership.companies
            ]
        }
      partnerships.append(partnership_data)
  
    return f"Found {len(partnerships)} partnerships", partnerships

@tool
def search_node (name: str, category: str) -> str:
  """Check if a node is already stored in the graph database."""

  data = load_json()

  search_prompt = f"""You are an expert search algorithm. I am going to give you the name of a company, and you have to check if the company is already stored in the list.

  You have 3 possibilities to answer:

  1. Exact match found: If a node with the same or very similar name already exists in the list:
    → Answer: "Node found with id [ID] and name [NAME]."

  2. Similar node found (requires merging): If a node represents the same company but with a different name variation (subsidiaries, regional divisions, different spelling):
    → Answer: "Similar node found with id [ID] and name [NAME]. Should be merged/renamed to [SUGGESTED_NAME] to avoid duplicates"
    
    Examples of nodes that should be merged:
    - "Toyota China" and "Toyota" → merge to "Toyota"
    - "Tesla Motors" and "Tesla Inc" → merge to "Tesla"
    - "Waymo" and "Google Self-Driving" → merge to "Waymo (Google)"
    - "General Motors" and "GM" → merge to "General Motors"

  3. No match found: If the company is genuinely new and doesn't match any existing node:
    → Answer: "Node not found. Need to create a new node with name [NAME] and category [CATEGORY]"

  Important guidelines:
  - Consider company subsidiaries, regional divisions, and alternative names as the same entity
  - Prioritize the most commonly used or official company name
  - When in doubt between similar companies, prefer merging over creating duplicates
  - Be case-insensitive in your search
  - Look for partial matches and common abbreviations

  The node to search: {name}
  The category of the node: {category}
  The existing nodes list: {data["nodes"]}

  Your analysis:
  """

  search_result = llm.invoke(search_prompt)

  return search_result.content

@tool
def add_node (name: str, category: str) -> tuple[str, dict]:
  """Add a node to the nodes list."""

  data = load_json()

  id = len(data["nodes"])
  new_node = {"id": id, "name": name, "category": category}
  data["nodes"].append(new_node)

  save_json(data)

  return f"Node added with the id {id}: ", new_node

@tool
def rename_node (id: int, new_name: str) -> str:
  """Update the name of a node from a given ID."""
  
  data = load_json()
  nodes = data["nodes"]

  if id >= len(nodes) or id < 0:
      return f"Error: Node with id {id} does not exist."
  
  previous_name = nodes[id]["name"]
  nodes[id]["name"] = new_name

  data["nodes"] = nodes
  save_json(data)

  return f"Node with id {id} renamed from '{previous_name}' to '{new_name}'."

@tool
def search_edge(from_id: int, to_id: int) -> str:
  """Check if an edge is already stored in the graph database."""

  data = load_json()

  for edge in data["edges"]:
    if (edge["from"] == from_id and edge["to"] == to_id) or (edge["from"] == to_id and edge["to"] == from_id):
      return f"Edge found between nodes {from_id} and {to_id}."
    
  return f"Edge not found. Need to create a new edge between the nodes {from_id} and {to_id}."

@tool
def add_edge (from_id: int, to_id: int) -> tuple[str, dict]:
  """Add an edge in the edges list."""

  data = load_json()
  edges = data["edges"]

  id = len(edges)
  new_edge = {"id": id, "from": from_id, "to": to_id}
  edges.append(new_edge)

  data["edges"] = edges
  save_json(data)

  return f"edge added with the id {id}: ", new_edge

# The agent
llm = init_chat_model("gpt-4o", model_provider="openai")

tools = [search_node, search_edge, add_node, rename_node, add_edge, extract_partnerships]

prompt = ChatPromptTemplate.from_messages([
    ("system", """You are an assistant specialized in analyzing partnerships in the self-driving car market.

Your capabilities:
- Extract partnerships from text articles
- Create and manage a node-based graph of partnerships
- Analyze companies and their relationships to avoid duplicates in the graph

Process for analyzing partnerships:
1. Use extract_partnerships to analyze the input text
2. For each company found, check if similar nodes already exist (search_node) or rename it if necessary (rename_node)
3. Add new nodes if necessary (add_node)
4. Check if relationships already exist (search_edge)
5. Add new edges between the nodes if necessary (add_edge)

Be precise and remember that the graph is undirected. Merge companies and their subsidiaries into a single node."""),
    
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}")
])

agent = create_tool_calling_agent(
  llm=llm,
  tools=tools,
  prompt=prompt
)

agent_executor = AgentExecutor(
  agent=agent,
  tools=tools,
  verbose=True,
  max_iterations=20,
  return_intermediate_steps=True
)

# Invoke the agent on one article
result = agent_executor.invoke({
  "input": news["content"][48]
})