import os
import re
import yaml
from rich import print
# from langchain import OpenAI
from langchain.chat_models import AzureChatOpenAI, ChatOpenAI

from LLMAgent.ConversationBot import ConversationBot

from LLMAgent.dataTools import (
    roadVolumeTrend,
    roadVolume,
    roadNameToID,
    plotGeoHeatmap,
    getCurrentTime,
    roadVisulization,
    odVolume,
    odMap
)

import gradio as gr
import openai.api_requestor
openai.api_requestor.TIMEOUT_SECS = 30

OPENAI_CONFIG = yaml.load(open('config.yaml'), Loader=yaml.FullLoader)
if OPENAI_CONFIG['OPENAI_API_TYPE'] == 'azure':
    os.environ["OPENAI_API_TYPE"] = OPENAI_CONFIG['OPENAI_API_TYPE']
    os.environ["OPENAI_API_VERSION"] = OPENAI_CONFIG['AZURE_API_VERSION']
    os.environ["OPENAI_API_BASE"] = OPENAI_CONFIG['AZURE_API_BASE']
    os.environ["OPENAI_API_KEY"] = OPENAI_CONFIG['AZURE_API_KEY']
    llm = AzureChatOpenAI(
        deployment_name=OPENAI_CONFIG['AZURE_MODEL'],
        temperature=0,
        max_tokens=1024,
        request_timeout=60
    )
elif OPENAI_CONFIG['OPENAI_API_TYPE'] == 'openai':
    os.environ["OPENAI_API_KEY"] = OPENAI_CONFIG['OPENAI_KEY']
    llm = ChatOpenAI(
        temperature=0,
        model_name='gpt-3.5-turbo-16k-0613',  # or any other model with 8k+ context
        max_tokens=1024,
        request_timeout=60
    )

if not os.path.exists('./fig/'):
    os.mkdir('./fig/')

toolModels = [
    roadVolumeTrend('./fig/'),
    roadVolume(),
    roadNameToID(),
    plotGeoHeatmap('./fig/'),
    getCurrentTime(),
    roadVisulization('./fig/'),
    odVolume(),
    odMap('./fig/')
]

botPrefix = """
# 1. You are a AI to assist human with traffic big-data analysis and visulization. #
# 2. You have access to the road network and traffic flow data in Xuancheng City, Anhui Province, China.
# 3. Whenever you are about to come up with a thought, recall the human message to check if you already have enough information for the final answer. If so, you shouldn't infer or fabricate any more needs or questions based on your own ideas.
# 4. You are forbidden to fabricate any tool names. If you can not find any appropriate tool for your task, try to do it using your own ability and knowledge as a chat AI.
# 5. Remember what tools you have used, DONOT use the same tool repeatedly. Try to use the least amount of tools.
# 6. DONOT fabricate any input parameters when calling tools! Check if you have the correct format of input parameters before calling tools!
# 7. When you encounter tabular content in Observation, make sure you output the tabular content in markdown format into your final answer.
# 8. Your tasks will be highly time sensitive. When generating your final answer, you need to state the time of the data you are using.
# 9. It's ok if the human message is not a traffic data related task, don't take any action and just respond to it like an ordinary conversation using your own ability and knowledge as a chat AI.
# 10. When you realize that you need to clarify what the human wants, end your actions and ask the human for more information as your final answer.
"""

bot = ConversationBot(llm, toolModels, botPrefix, verbose=True)


def reset(chat_history: list, thoughts: str):
    chat_history = []
    thoughts = ""
    bot.agent_memory.clear()
    bot.ch.memory = [[]]
    return chat_history, thoughts


def respond(msg: str, chat_history: list, thoughts: str):
    res, cb = bot.dialogue(msg)
    regex = re.compile(r'`([^`]+)`')
    try:
        filenames = regex.findall(res)
    except AttributeError:
        filenames = None
    if filenames:
        chat_history += [(msg, None)]
        for fn in filenames:
            chat_history += [(None, (fn,))]
        chat_history += [(None, res)]
    else:
        chat_history += [(msg, res)]

    thoughts += f"\n>>> {msg}\n"
    for actionMemory in bot.ch.memory[-2]:
        thoughts += actionMemory
        thoughts += '\n'
    thoughts += f"<<< {res}\n"
    return "", chat_history, thoughts


with gr.Blocks(
    title="Traffic Data Process Bot", theme=gr.themes.Base(text_size=gr.themes.sizes.text_md)
) as demo:
    with gr.Row(visible=True, variant="panel"):
        with gr.Column(visible=True, variant='default'):
            chatbot = gr.Chatbot(scale=2, height=650)

            with gr.Row():
                humanMsg = gr.Textbox(scale=2)
                submitBtn = gr.Button("Submit", scale=1)
            clearBtn = gr.ClearButton()
            gr.Examples(
                label='You may want to ask the following questions:',
                examples=[
                    "Show me the OD map from 7am to 9am today.",
                    "Show me the current network heatmap.",
                    "Show me the traffic volume of OD pairs from 5pm to 7pm yesterday.",
                    "Show me the traffic volume data overview of yesterday in a table.",
                    "青弋江西大道在哪？",
                    "How's the traffic volume trend of road 1131 yesterday?"
                ],
                inputs=[humanMsg],
                # outputs=[humanMsg, chatbot],
                # fn=testFunc
            )
        ReActMsg = gr.Text(
            label="Thoughts and Actions of the Chatbot",
            interactive=False,
            lines=50
        )

    humanMsg.submit(
        respond,
        [humanMsg, chatbot, ReActMsg],
        [humanMsg, chatbot, ReActMsg]
    )
    submitBtn.click(
        respond,
        [humanMsg, chatbot, ReActMsg],
        [humanMsg, chatbot, ReActMsg]
    )
    clearBtn.click(reset, [chatbot, ReActMsg], [chatbot, ReActMsg])

if __name__ == "__main__":
    demo.launch()

DataProcessBot.py
lijlansg/TrafficGPT
from rich import print
from typing import Any, List
from langchain import LLMChain
from langchain.agents import Tool
from langchain.chat_models import AzureChatOpenAI
from LLMAgent.callbackHandler import CustomHandler
from langchain.callbacks import get_openai_callback
from langchain.memory import ConversationBufferMemory
from langchain.agents import ZeroShotAgent, Tool, AgentExecutor


temp = """
You need to recall the original 'Question' before comming up with a 'Thought'. 
2. You need to determine whether the human message is a traffic simulation control command or a question before making any move. If it is a traffic simulation control command, just execute the command and don't do any further information analysis. If it's neither, try to respond to it using your own ability and knowledge as a chat AI
5. Stop calling other tools if you have enough information to answer the questions or already fulfilled the commands explicitly mentioned in the human message. 
"""

# prefix = """
# 1. You are a AI to assist human with traffic simulation control or making traffic and transportation decisions.
# 2. You need to determine whether the human message is a traffic simulation control command or a question before making any move. If it is a traffic simulation control command, just execute the command and don't do any further information analysis. If
# 3. You need to remeber the human message exactly. Your only purpose is to complete the task that is explicitly expressed in the human message.
# 4. Whenever you are about to come up with a thought, recall the human message to check if you already have enough information for the final answer. If so, you shouldn't infer or fabricate any more needs or questions based on your own ideas.
# 5. You are forbidden to fabricate any tool names. If you can not find any appropriate tool for your task, try to do it using your own ability and knowledge as a chat AI.
# 6. Remember what tools you have used, DONOT use the same tool repeatedly. Try to use the least amount of tools.
# 7. DONOT fabricate any input parameters when calling tools!
# 8. When you encounter tabular content in Observation, make sure you output the tabular content in markdown format into your final answer.
# 9. When you realize that existing tools are not solving the problem at hand, you need to end your actions and ask the human for more information as your final answer.
# """

# simbot+report

# prefix = """
# [WHO ARE YOU]
# You are a AI to assist human with traffic simulation control, making traffic and transportation decisions, or providing traffic analysis reports. Although you have access to a set of tools, your abilities are not limited to the tools at your disposal
# [YOUR ACTION GUIDLINES]
# 1. You need to determine whether the human message is a traffic simulation control command or a question before making any move. If it is a traffic simulation control command, just execute the command and don't do any further information analysis. If
# 2. You need to remeber the human message exactly. Your only purpose is to complete the task that is explicitly expressed in the human message.
# 3. Whenever you are about to come up with a thought, recall the human message to check if you already have enough information for the final answer. If so, you shouldn't infer or fabricate any more needs or questions based on your own ideas.
# 4. Remember what tools you have used, DONOT use the same tool repeatedly. Try to use the least amount of tools.
# 5. If you can not find any appropriate tool for your task, try to do it using your own ability and knowledge as a chat AI.
# 6. When you encounter tabular content in Observation, make sure you output the tabular content in markdown format into your final answer.
# 7. When you realize that existing tools are not solving the problem at hand, you need to end your actions and ask the human for more information as your final answer.
# [THINGS YOU CANNOT DO]
# You are forbidden to fabricate any tool names.
# You are forbidden to fabricate any input parameters when calling tools!
# [HOW TO GENERATE TRAFFIC REPORTS]
# Act as a human. And provide as much information as possible, including file path and tabular datasets.
# When human need to provede a report of the traffic situation of a road network, they usually start by observing the operation of the network,
# find a few intersections in the network that are in a poor operating condition, as well as their locations, try to optimize them,
# and evaluate which parameters have become better and which ones are worse after the optimization. And form a report of the complete thought process in markdown format.
# For example:
# Macroscopic traffic operations on the entire road network can be viewed on the basis of road network heatmaps: 'replace the correct filepath here'.
# To be more specific, these 5 intersections are in the worst operation status.
# |    |   Juction_id |   speed_avg |   volume_avg |   timeLoss_avg |
# |---:|-------------:|------------:|-------------:|---------------:|
# |  0 |         4605 |     8.02561 |      734.58 |        8155.83 |
# |  1 |         4471 |     8.11299 |      797.92 |       16500.6  |
# |  2 |         4493 |     8.36199 |      532.26 |        8801.71 |
# |  3 |         4616 |     8.62853 |      898.08 |        5897.33 |
# |  4 |         4645 |     9.38659 |      360.03 |       11689    |
# the locations of these intersections are shown in the map: 'replace the correct filepath here'.
# I tried to optimize the traffic signal shceme of them and run the simulation again.
# The new traffic stauts of these 5 intersections are as follows:
# |    |   Juction_id |   speed_avg |   volume_avg |   timeLoss_avg |
# |---:|-------------:|------------:|-------------:|---------------:|
# |  0 |         4605 |     5.02561 |      1734.58 |        9155.83 |
# |  1 |         4471 |     5.11299 |      1797.92 |       17500.6  |
# |  2 |         4493 |     5.36199 |      1532.26 |        9901.71 |
# |  3 |         4616 |     5.62853 |      1898.08 |        6897.33 |
# |  4 |         4645 |     5.38659 |      1360.03 |       13689    |
# According to the data above, after optimization, Traffic volume has increased at these intersections, but average speeds have slowed and time loss have become greater.
# """

suffix = """Begin!"

{chat_history}
Question: {input}
{agent_scratchpad}"""


class ConversationBot:
    def __init__(
            self, llm: AzureChatOpenAI, toolModels: List,
            customedPrefix: str, verbose: bool = False
    ) -> Any:
        self.ch = CustomHandler()
        tools = []

        for ins in toolModels:
            func = getattr(ins, 'inference')
            tools.append(
                Tool(
                    name=func.name,
                    description=func.description,
                    func=func
                )
            )

        prompt = ZeroShotAgent.create_prompt(
            tools,
            prefix=customedPrefix,
            suffix=suffix,
            input_variables=["input", "chat_history", "agent_scratchpad"],
        )
        self.agent_memory = ConversationBufferMemory(memory_key="chat_history")

        llm_chain = LLMChain(llm=llm, prompt=prompt)
        agent = ZeroShotAgent(
            llm_chain=llm_chain,
            tools=tools, verbose=verbose
        )
        self.agent_chain = AgentExecutor.from_agent_and_tools(
            agent=agent, tools=tools,
            verbose=verbose, memory=self.agent_memory,
            handle_parsing_errors="Use the LLM output directly as your final answer!"
        )

    def dialogue(self, input: str):
        print('TransGPT is running, Please wait for a moment...')
        with get_openai_callback() as cb:
            res = self.agent_chain.run(input=input, callbacks=[self.ch])
        # print('History: ', self.agent_memory.buffer)
        return res, cb

LLMAgent/ConversationBot.py
lijlansg/TrafficGPT
from matplotlib import pyplot as plt
from xml.etree import ElementTree as ET
import matplotlib
matplotlib.use('TkAgg')


class Lane:
    def __init__(
            self, id: str, speed: float, length: float,
            raw_shape: str
    ) -> None:
        self.id = id
        self.speed = speed
        self.length = length
        self.shape = self.process_raw_shape(raw_shape)

    def process_raw_shape(self, raw_shape: str) -> [[float]]:
        raw_list = raw_shape.split(' ')
        shape = [list(map(float, p.split(','))) for p in raw_list]
        return shape

    def plot_self(self, color: str, alpha: float):
        plt.plot(
            list(zip(*self.shape))[0],
            list(zip(*self.shape))[1],
            color=color, alpha=alpha
        )


class Edge:
    def __init__(self, id: str, fromNode: str, toNode: str) -> None:
        self.id = id,
        self.fromNode = fromNode
        self.toNode = toNode
        self.lanes: dict[str, Lane] = {}
        self.next_edges: set[str] = set()
        self.length: float = 0
        self.freeFlowSpeed = 13.89

    def get_length(self) -> None:
        max_length = 0
        for lane in self.lanes.values():
            if lane.length > max_length:
                max_length = lane.length
        self.length = max_length
        self.capacity = 1800.0 * len(self.lanes)

    def plot_self(self, color: str, alpha: float) -> None:
        for lane in self.lanes.values():
            lane.plot_self(color, alpha)


class Junction:
    def __init__(self, id: str) -> None:
        self.id = id
        self.inEdges: set[Edge] = set()
        self.outEdges: set[Edge] = set()

    def calCap(self) -> float:
        jCap = 0
        nEdges = len(self.inEdges)
        for edge in self.inEdges:
            jCap += edge.capacity / nEdges

        return jCap


class Graph:
    def __init__(self) -> None:
        self.edges: dict[str, Edge] = {}
        self.junctions: dict[str, Junction] = {}

    def get_junction(self, jid: str):
        return self.junctions[jid]

    @property
    def validJunctions(self):
        invalidJunctions = set(
            ['9173', '4680', '4562', '4423', '4609']
        )
        return self.junctions.keys() - invalidJunctions

    def get_edge(self, eid: str):
        return self.edges[eid]

    def getEdgeByJunction(self, fnode: str, tnode: str):
        fj = self.get_junction(fnode)
        for edge in fj.outEdges:
            if edge.toNode == tnode:
                return edge
        raise ValueError(f'There is no edge between {fnode} and {tnode}.')

    def plot_self(self):
        for edge in self.edges.values():
            edge.plot_self('grey', 0.3)


def completeJunctions(graph: Graph):
    for ek, ev in graph.edges.items():
        fj = graph.get_junction(ev.fromNode)
        tf = graph.get_junction(ev.toNode)
        fj.outEdges.add(ev)
        tf.inEdges.add(ev)


def build_graph(file: str) -> Graph:
    graph = Graph()
    eTree = ET.parse(file)
    root = eTree.getroot()
    for child in root:
        if child.tag == 'edge':
            eid = child.attrib['id']
            if ':' not in eid:
                fromNode = child.attrib['from']
                toNode = child.attrib['to']
                edge = Edge(eid, fromNode, toNode)
                for gchild in child:
                    lid = gchild.attrib['id']
                    speed = float(gchild.attrib['speed'])
                    length = float(gchild.attrib['length'])
                    raw_shape = gchild.attrib['shape']
                    lane = Lane(
                        lid, speed, length, raw_shape)
                    edge.lanes[lid] = lane
                    edge.get_length()
                graph.edges[eid] = edge
            else:
                continue
        elif child.tag == 'junction':
            jid = child.attrib['id']
            if ':' not in jid:
                graph.junctions[jid] = Junction(jid)
        elif child.tag == 'connection':
            from_edge_id = child.attrib['from']
            to_edge_id = child.attrib['to']
            if ':' not in from_edge_id and ':' not in to_edge_id:
                from_edge = graph.get_edge(from_edge_id)
                from_edge.next_edges.add(to_edge_id)
            else:
                continue
        else:
            continue

    completeJunctions(graph)
    return graph

LLMAgent/buildGraph.py
lijlansg/TrafficGPT
from __future__ import annotations
from typing import Any, Dict, List, Optional
from uuid import UUID
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import AgentAction, AgentFinish
from rich import print


class CustomHandler(BaseCallbackHandler):
    def __init__(self) -> None:
        super().__init__()
        self.memory: List[List[str]] = [[]]

    def on_agent_finish(self, finish: AgentFinish, *, run_id: UUID, parent_run_id: UUID | None = None, **kwargs: Any) -> Any:
        self.memory.append([])
        return super().on_agent_finish(finish, run_id=run_id, parent_run_id=parent_run_id, **kwargs)

    def on_agent_action(self, action: AgentAction, *, run_id: UUID, parent_run_id: UUID | None = None, **kwargs: Any) -> Any:
        self.memory[-1].append(action.log)
        return super().on_agent_action(action, run_id=run_id, parent_run_id=parent_run_id, **kwargs)

LLMAgent/callbackHandler.py
lijlansg/TrafficGPT
import pandas as pd
import matplotlib.pyplot as plt

from LLMAgent.dbConnector import fetch_from_database
from LLMAgent.plotGeoMap import plot_geo_heatmap, plot_road_segements, plot_OD_map
from LLMAgent.getTime import get_time_period, get_fake_current_time


def prompts(name, description):
    def decorator(func):
        func.name = name
        func.description = description
        return func

    return decorator



class roadVolumeTrend:
    def __init__(self, figfolder: str) -> None:
        self.figfolder = figfolder

    @prompts(name='Visualize Road Traffic Trends',
             description="""
            This tool is used to show the traffic volume trend of several target roads on a specific day by visualize it on a histogram.
            Use this tool more than others if the question is about traffic trends.
            The output will tell you whether you have finished this command successfully.
            The input should be one sigle semicolon-separated string that separate two parts of information, and the two parts to the left and right of the semicolon are date and roadID, respectively. 
            The input date should be an integer, representing the day of the month. The input roadID should be a "|" seperated string, representing the id of the target roads. For example: 15;1076|30188""")
    def inference(self, target: str) -> str:

        day, road_id = target.replace(' ', '').split(';')

        # query
        query = f"""
        select date_trunc('hour', departure_time) as hours, count(*) as volume  
        from the_synthetic_individual_level_trip_dataset
        where path similar to '%({road_id})%' AND extract(day from departure_time) = {day}
        group by hours
        order by hours;
        """

        rows = fetch_from_database(query)

        time = []
        volume = []
        for row in rows:
            time.append(row[0].strftime('%Y-%m-%d %H'))
            volume.append(row[-1])


        plt.figure(figsize=(16,8),dpi=300)
        plt.bar(time,volume)
        plt.xticks(rotation=70)
        fig_path = f'{self.figfolder}volume_bar.png'
        plt.savefig(fig_path, dpi=300)

        return f"You have successfully visualized the Road Traffic Trends. And your final answer should include this sentence without changing anything: The histogram showing the traffic trend is kept at: `{fig_path}`."


class roadVolume:
    def __init__(self) -> None:
        pass

    @prompts(name='Get Road Volume',
             description="""
            This tool is used to count the traffic volume of roads during a specific time period.
            Use this tool more than others if the question is about specific traffic volume.
            The output will provid road volume information in a tabular dataset.
            The input should be one sigle semicolon-separated string that separate two parts of information, and the two parts to the left and right of the semicolon are time_period and targets, respectively: 
            1. The time_period: The input time_period should be a comma seperated string, with each part representing the start and end date and time of the time_period in the "YYYY-MM-DD HH:MM:SS" format.
            2. The targets: If you do not have information of any specific road ID, the input should be a string: 'None', and the tool will output an overview data for the final answer. If you have specific target road IDs, the input should be a comma seperated string, with each part representing a target road ID. Only if you can find a word 'all' in the human message, the input can be a string: 'All'.
            For example: 2019-08-13 08:00:00,2019-08-13 09:00:00;None  /  2019-08-13 08:00:00,2019-08-13 09:00:00;2023,3486""")
    def inference(self, inputs: str) -> str:
        
        time_period, target = inputs.split(';')
        begin, end = time_period.split(',')
        if 'None' in target or 'All' in target:
            # print('no target' + target.replace(' ', ''))
            have_target = False
            target_road_id = []
        else:
            have_target = True
            target_road_id = target.split(',')
            # print('target'+ str(target.replace(' ', '').split(','))) 
            idlist = ', '.join(f"'{item}'" for item in target_road_id)
        
        # query
        if have_target == False and 'None' in target:
            query = f"""
            select road as road_id, SUM(road_count) as volume
            from road_volume_per_hour
            WHERE hour_start >= '{begin}' AND hour_start < '{end}'
            group by road_id
            ORDER BY volume DESC
            LIMIT 5;
            """
        elif have_target == False and 'All' in target:
            query = f"""
            select road as road_id, SUM(road_count) as volume
            from road_volume_per_hour
            WHERE hour_start >= '{begin}' AND hour_start < '{end}'
            group by road_id
            ORDER BY volume DESC
            """
        else:
            query = f"""
            select road as road_id, SUM(road_count) as volume
            from road_volume_per_hour
            WHERE hour_start >= '{begin}' AND hour_start < '{end}' AND road IN ({idlist})
            group by road_id
            ORDER BY volume DESC
            """

        rows = fetch_from_database(query)
        data = pd.DataFrame(rows,columns=['road_id','volume'])

        if 'None' in target:
            msg = 'No specific target roads. The human user just wants to see an overview. So, I can show you the overview by providing traffic data for the 5 highest volume roads by default. Make sure you output the tabular content in markdown format into your final answer. \n'
            return msg + data.to_markdown()
        elif 'All' in target:
            msg = 'Here are the traffic volume of all roads. Make sure you output the tabular content in markdown format into your final answer. \n'
            return msg + data.to_markdown()
        else:
            msg = 'Here are the traffic status of your targeted roads. Make sure you output the tabular content in markdown format into your final answer. \n'
            return msg + data.to_markdown()


class roadNameToID:
    def __init__(self) -> None:
        pass

    @prompts(name='Get Road_ID From Road Name',
             description="""
            This tool is used to get the specific road id according to road names.
            Use this tool before others if the target road info given by human user is the Name of Road rather than ids(which is series of figures). In this case, use this tool to get the road id as input for other tools.
            The output will tell you the id of the road that corresponds to the road name in a tabular dataset.
            The input should be a string representing the road name, for example: 青弋江西大道.
            NOTE that a road name may correspond to several road ids, remember them all! """)
    def inference(self, target: str) -> str:

        road_name = target.replace(' ', '')

        # query
        query = f"""
        select objectid,mc
        from topo_centerroad
        where mc = '{road_name}'
        """

        rows = fetch_from_database(query)
        data = pd.DataFrame(rows,columns=['road_id','road_name'])
        road_id_list = data['road_id'].tolist()
        road_ids = ', '.join(str(int(road)) for road in road_id_list)
        msg = f'Here are the road ids of all roads corresponding to the road name {road_name}:'
        return msg + road_ids


class plotGeoHeatmap:
    def __init__(self, figfolder: str) -> None:
        self.figfolder = figfolder

    @prompts(name='Plot Heatmap',
             description="""
            This tool is used to show the traffic status of the network by ploting a heatmap of the road network at a certain point in time.
            If you do not have a specific point in time, use "Get Current Time" tool to get the time before using this tool.
            Use this tool more than others if the human user want to see the traffic status of the network.
            The output will tell you the file path of the heat mapas a supplementary information for you to provide the final answer. 
            The input must be the target point in time in the "YYYY-MM-DD HH:MM:SS" format. 
            
            """)
    def inference(self, time: str) -> str:
        try:
            start, end = get_time_period(time)
        except:
            return "Wrong format of input parameters, The input must be the target point in time in the 'YYYY-MM-DD HH:MM:SS' format."

        # query
        query = f"""SELECT mc, objectid, volume, geom 
        FROM topo_centerroad join (
        select road as road_ID, SUM(road_count) as volume
        from road_volume_per_hour
        WHERE hour_start >= '{start}' AND hour_start < '{end}'
        group by road_ID) as traffic_volume
        on CAST(ROUND(objectid) as varchar) = traffic_volume.road_ID;"""

        fig_path = plot_geo_heatmap(query,self.figfolder)

        return f"The heat map is ploted according to traffic volume data at {time}. And your final answer should include this sentence without changing anything: the road network heat map is kept at: `{fig_path}`."
    

class getCurrentTime:
    def __init__(self) -> None:
        pass

    @prompts(name='Get Current Time',
             description="""
            This tool is used to get the current date and time.
            Use this tool before others if you need to know the current date and time.
            The output will tell you the current time in the "YYYY-MM-DD HH:MM:SS" format.
            """)
    def inference(self, time: str) -> str:

        # 根据数据库数据，获取虚拟的当前时间
        current_time = get_fake_current_time()

        return f"The current time is {current_time}"
    

class roadVisulization:
    def __init__(self, figfolder: str) -> None:
        self.figfolder = figfolder

    @prompts(name='Visualize Roads',
             description="""
            This tool is used to show the locations of several target roads by visualize them on a map.
            Use this tool more than others if the question is about locations of roads.
            The output will tell you whether you have finished this command successfully.
            The input should be a comma seperated string, with each part be a series of figures representing a target road_id. 
            For example: 1076,30188 """)
    def inference(self, target: str) -> str:

        road_ids = target.replace(' ', '')

        fig_path = plot_road_segements(road_ids, self.figfolder)

        return f"You have successfully visualized the location of road {target} on the following map. And your final answer should include the following sentence without changing words: The location of road {target} is kept at: `{fig_path}`."
    

class odVolume:
    def __init__(self) -> None:
        pass

    @prompts(name='Get OD Volume',
             description="""
            This tool is used to count the traffic volume of OD pairs during a specific time period.
            Use this tool more than others if the question is about OD pair traffic volume.
            The output will provid OD pair volume information in a tabular dataset.
            The input should be one sigle semicolon-separated string that separate two parts of information, and the two parts to the left and right of the semicolon are time_period and Number of data items you want to see, respectively: 
            1. The time_period: The input time_period should be a comma seperated string, with each part representing the start and end date and time of the time_period in the "YYYY-MM-DD HH:MM:SS" format.
            2. Number of data items: If you want to see the traffic volume data of OD pairs with top 10 traffic volume, the input should be 10. If no specific number of data items given, the input should be 5 as default.
            For example: 2019-08-13 08:00:00,2019-08-13 09:00:00;5 """)
    def inference(self, inputs: str) -> str:
        time_period, number = inputs.split(';')
        begin, end = time_period.split(',')
        if number == "None":
            N =5
        else:
            N = eval(number)

        query = f"""
        SELECT o_zone, d_zone, COUNT(*) AS od_pair_volume
        FROM the_synthetic_individual_level_trip_dataset
        WHERE departure_time >= '{begin}' AND departure_time < '{end}'
        GROUP BY o_zone, d_zone
        ORDER BY od_pair_volume DESC
        LIMIT {N};"""

        rows = fetch_from_database(query)
        data = pd.DataFrame(rows,columns=['o_zone','d_zone','od_pair_volume'])

        msg = f'Here are the traffic volume data of top{N} OD pairs. Make sure you output the tabular content in markdown format into your final answer. \n'
        return msg + data.to_markdown()
    

class odMap:
    def __init__(self, figfolder: str) -> None:
        self.figfolder = figfolder

    @prompts(name='Plot OD Map',
             description="""
            This tool is used to visulize the traffic volume of OD pairs during a specific time period on a OD map.
            Use this tool more than others if the question is about OD map.
            The output will tell you the file path of the OD map as a supplementary information for you to provide the final answer. 
            The input time_period should be a comma seperated string, with each part representing the start and end date and time of the time_period in the "YYYY-MM-DD HH:MM:SS" format.
            
            """)
    def inference(self, time: str) -> str:

        begin, end = time.split(',')
        
        fig_path = plot_OD_map(begin, end, self.figfolder)

        return f"The OD map is ploted according to traffic volume data at from {begin} to {end}. And your final answer should include this sentence without changing anything: The OD map is kept at: `{fig_path}`."
    
LLMAgent/dataTools.py
lijlansg/TrafficGPT
import sqlalchemy as sa
import yaml

def fetch_from_database(query):

    # 读取配置文件
    with open("dbconfig.yaml", "r") as f:
        config = yaml.safe_load(f)

    # 连接到PostgreSQL数据库
    db_uri = f"postgresql://{config['username']}:{config['password']}@{config['host']}:{config['port']}/{config['db_name']}"
    engine = sa.create_engine(db_uri)
    conn = engine.connect()
    
    # 执行查询
    query_obj = sa.text(query)
    result = conn.execute(query_obj)
    conn.close()
    return result

LLMAgent/dbConnector.py
lijlansg/TrafficGPT
from datetime import datetime, timedelta

def get_fake_current_time():
    # 获取当前时间
    current_time = datetime.now()

    # 修改日期为2019年8月16日
    new_date = datetime(2019, 8, 16, current_time.hour, current_time.minute, current_time.second)

    # 格式化时间为指定格式
    formatted_time = new_date.strftime('%Y-%m-%d %H:%M:%S')

    return formatted_time

def get_time_period(given_time_str):
    # 将字符串转换为datetime对象
    given_time = datetime.strptime(given_time_str, '%Y-%m-%d %H:%M:%S')

    # 获取该时间所在的小时
    hour = given_time.hour

    # 构造小时区间的起始时间和结束时间
    start_hour = datetime(given_time.year, given_time.month, given_time.day, hour, 0, 0)
    end_hour = start_hour + timedelta(hours=1)

    return start_hour, end_hour
LLMAgent/getTime.py
lijlansg/TrafficGPT
import geopandas as gpd
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.colorbar import ColorbarBase
import sqlalchemy as sa
import yaml


def plot_geo_heatmap(query,figfolder):
    
    # 读取配置文件
    with open("dbconfig.yaml", "r") as f:
        config = yaml.safe_load(f)

    # 连接到PostgreSQL数据库
    db_uri = f"postgresql://{config['username']}:{config['password']}@{config['host']}:{config['port']}/{config['db_name']}"
    engine = sa.create_engine(db_uri)
    conn = engine.connect()
    query_obj = sa.text(query)
    
    # 获取几何数据和交通流量数据
    gdf = gpd.GeoDataFrame.from_postgis(query_obj, conn, geom_col='geom')

    # 关闭数据库连接
    conn.close()

    # 创建渐变颜色映射和线条宽度映射
    colors = [(0, 1, 0), (1, 1, 0), (1, 0, 0)]  # 绿-黄-红的渐变色
    cmap = LinearSegmentedColormap.from_list('custom_cmap', colors, N=100)
    volume_normalized = (gdf['volume'] - gdf['volume'].min()) / (gdf['volume'].max() - gdf['volume'].min())

    # 设置绘图参数
    fig, ax = plt.subplots(figsize=(8, 6), dpi=800)
    ax.set_aspect('equal')

    # 绘制地图，根据交通流量数据设置颜色和线条宽度
    gdf['geom'].plot(ax=ax, cmap=cmap, linewidth=volume_normalized * 5)

    # 添加颜色图例（放在左侧）
    cax_position = [0.08, 0.25, 0.02, 0.5]
    norm = plt.Normalize(gdf['volume'].min(), gdf['volume'].max())
    cax = fig.add_axes(cax_position)  # 调整位置参数
    cb = ColorbarBase(cax, cmap=cmap, norm=norm, orientation='vertical')

    # 添加标题
    cb_title = 'Traffic Volume'
    cb_title_fontsize = 12
    plt.text(cax_position[0] + cax_position[2] / 2 + 0.01, cax_position[1] + cax_position[3] + 0.01,
            cb_title, fontsize=cb_title_fontsize, ha='center', va='bottom', fontname='Times New Roman', weight='bold', transform=fig.transFigure)

    # 隐藏坐标轴
    ax.set_axis_off()

    # 输出热力图
    fig_path = f'{figfolder}heatmap.png'
    fig.savefig(fig_path, dpi=800)
    return fig_path


def plot_road_segements(road_ids,figfolder):
    
    # 读取配置文件
    with open("dbconfig.yaml", "r") as f:
        config = yaml.safe_load(f)

    # 连接到PostgreSQL数据库
    db_uri = f"postgresql://{config['username']}:{config['password']}@{config['host']}:{config['port']}/{config['db_name']}"
    engine = sa.create_engine(db_uri)
    conn = engine.connect()

    # 获取几何数据和路名id
    query1 = "SELECT mc, objectid, geom FROM topo_centerroad"
    query1_obj = sa.text(query1)
    gdf1 = gpd.GeoDataFrame.from_postgis(query1_obj, conn, geom_col='geom')

    query2 = f"SELECT mc, objectid, geom FROM topo_centerroad WHERE objectid IN ({road_ids})"
    query2_obj = sa.text(query2)
    gdf2 = gpd.GeoDataFrame.from_postgis(query2_obj, conn, geom_col='geom')

    # 关闭数据库连接
    conn.close()

    # 设置参数
    annotation_params = {'fontsize': 8, 'color': 'black', 'fontname':'Times New Roman', 'weight':'bold'}  # 标注文本的参数
    fig, ax = plt.subplots(figsize=(8, 6), dpi=800)

    # 绘制地图，根据颜色数据设置颜色
    ax = gdf1['geom'].plot(ax=ax,color= 'grey')
    ax = gdf2['geom'].plot(ax=ax,color= 'red')

    # 标注每条路的id
    for idx, row in gdf2.iterrows():
        ax.text(row['geom'].centroid.x, row['geom'].centroid.y, int(row['objectid']), **annotation_params)

    fig_path = f'{figfolder}roads.png'
    fig.savefig(fig_path, dpi=800)
    return fig_path

def plot_OD_map(begin, end, figfolder):

    # 读取配置文件
    with open("dbconfig.yaml", "r") as f:
        config = yaml.safe_load(f)

    # 连接到PostgreSQL数据库
    db_uri = f"postgresql://{config['username']}:{config['password']}@{config['host']}:{config['port']}/{config['db_name']}"
    engine = sa.create_engine(db_uri)
    conn = engine.connect()

    # 获取OD数据
    query1 = f"""
    SELECT o_zone, d_zone, od_pair_count, ST_MakeLine(ST_MakePoint(ozone_long, ozone_lat), ST_MakePoint(dzone_long, dzone_lat)) AS geom 
    from (
            SELECT
                o_zone,
                d_zone,
                od_pair_count,
                ozone_info.longitude AS ozone_long,
                ozone_info.latitude AS ozone_lat,
                dzone_info.longitude AS dzone_long,
                dzone_info.latitude AS dzone_lat
            FROM
                (
                    SELECT
                        o_zone,
                        d_zone,
                        od_pair_count
                    FROM
                        (SELECT o_zone, d_zone, COUNT(*) AS od_pair_count
                                    FROM the_synthetic_individual_level_trip_dataset
                                    WHERE departure_time >= '{begin}' AND departure_time < '{end}'
                                    GROUP BY o_zone, d_zone
                                    ORDER BY od_pair_count DESC) as od_pair_count
                ) AS od
            JOIN
                zone_roads AS ozone_info ON od.o_zone = ozone_info.zone_id
            JOIN
                zone_roads AS dzone_info ON od.d_zone = dzone_info.zone_id
            ) as OD_volume;
    """
    query1_obj = sa.text(query1)
    gdf1 = gpd.GeoDataFrame.from_postgis(query1_obj, conn, geom_col='geom')

    # 获取地图数据
    query2 = "SELECT mc, objectid, geom FROM topo_centerroad"
    query2_obj = sa.text(query2)
    gdf2 = gpd.GeoDataFrame.from_postgis(query2_obj, conn, geom_col='geom')

    # 关闭数据库连接
    conn.close()

    # 创建渐变颜色映射（绿-黄-红）
    colors = [(0, 1, 0), (1, 1, 0), (1, 0, 0)]  # 绿-黄-红的渐变色
    cmap = LinearSegmentedColormap.from_list('custom_cmap', colors, N=100)
    volume_normalized = (gdf1['od_pair_count'] - gdf1['od_pair_count'].min()) / (gdf1['od_pair_count'].max() - gdf1['od_pair_count'].min())

    # 设置绘图参数
    fig, ax = plt.subplots(figsize=(8, 6), dpi=800)
    ax.set_aspect('equal')

    # 绘制地图，根据交通流量数据设置颜色和线条宽度
    gdf1['geom'].plot(ax=ax, cmap=cmap, linewidth=volume_normalized*5)
    gdf2['geom'].plot(ax=ax,color= 'grey')

    cax_position = [0.08, 0.25, 0.02, 0.5]
    norm = plt.Normalize(gdf1['od_pair_count'].min(), gdf1['od_pair_count'].max())
    cax = fig.add_axes(cax_position)  # 调整位置参数
    cb = ColorbarBase(cax, cmap=cmap, norm=norm, orientation='vertical')
    # cb.set_label('Traffic Volume', rotation=270, labelpad=15)  # 设置颜色图例标题

    # 添加标题
    cb_title = 'Traffic Volume'
    cb_title_fontsize = 12
    plt.text(cax_position[0] + cax_position[2] / 2 + 0.01, cax_position[1] + cax_position[3] + 0.01,
            cb_title, fontsize=cb_title_fontsize, ha='center', va='bottom', fontname='Times New Roman', weight='bold', transform=fig.transFigure)

    ax.set_axis_off()

    fig_path = f'{figfolder}ODmap.png'
    fig.savefig(fig_path, dpi=800)
    return fig_path

LLMAgent/plotGeoMap.py
lijlansg/TrafficGPT
from __future__ import absolute_import
from __future__ import print_function

import os
import sys
from xml.sax.handler import ContentHandler

import matplotlib
matplotlib.use('Agg')

import sumolib  # noqa
from sumolib.visualization import helpers  # noqa
import matplotlib.pyplot as plt  # noqa
import matplotlib  # noqa

class WeightsReader(ContentHandler):

    """Reads the dump file"""

    def __init__(self, value):
        super().__init__()
        self._edge2value = {}
        self._value = value
        self._intervals = []

    def startElement(self, name, attrs):
        if name == 'interval':
            self._time = float(attrs['begin'])
            self._edge2value[self._time] = {}
            self._intervals.append(self._time)

        elif name == 'edge':
            id = attrs['id']
            if self._value in attrs:
                self._edge2value[self._time][id] = float(
                    attrs[self._value])

def plot_heatmap(figfolder, args=None) -> str:
    """The main function; parses options and plots"""
    # ---------- build and read options ----------
    ap = sumolib.options.ArgumentParser()
    ap.add_option("-n", "--net", dest="net", metavar="FILE", required=True, type=ap.net_file,
                    help="Defines the network to read")
    ap.add_option("-i", "--dump-inputs", dest="dumps", metavar="FILE", required=True, type=ap.edgedata_file,
                    help="Defines the dump-output files to use as input")
    ap.add_option("-m", "--measures", dest="measures",
                    default="speed,entered", help="Define which measure to plot")
    ap.add_option("--min-width", dest="minWidth", type=float, default=.5,
                    help="Defines the minimum edge width")
    ap.add_option("--max-width", dest="maxWidth",
                    type=float, default=3, help="Defines the maximum edge width")
    ap.add_option("--log-colors", dest="logColors", action="store_true",
                    default=False, help="If set, colors are log-scaled")
    ap.add_option("--log-widths", dest="logWidths", action="store_true",
                    default=False, help="If set, widths are log-scaled")
    ap.add_option("--min-color-value", dest="colorMin",
                    type=float, default=None,
                    help="If set, defines the minimum edge color value")
    ap.add_option("--max-color-value", dest="colorMax",
                    type=float, default=None,
                    help="If set, defines the maximum edge color value")
    ap.add_option("--min-width-value", dest="widthMin",
                    type=float, default=None,
                    help="If set, defines the minimum edge width value")
    ap.add_option("--max-width-value", dest="widthMax",
                    type=float, default=None,
                    help="If set, defines the maximum edge width value")
    ap.add_option("-v", "--verbose", dest="verbose", action="store_true",
                    default=False,
                    help="If set, the script says what it's doing")
    ap.add_option("--color-bar-label", dest="colorBarLabel", default="",
                    help="The label to put on the color bar")
    ap.add_option("--internal", action="store_true",
                    default=False, help="include internal edges in generated shapes")

    # standard plot options
    matplotlib.use('Agg')
    helpers.addInteractionOptions(ap)
    helpers.addPlotOptions(ap)
    helpers.addNetOptions(ap)

    # Override the help string for the output option
    outputOpt = ap.get_option("output")
    outputOpt.help = ("Comma separated list of filename(s) the figure shall be written to; " +
                        "for multiple time intervals use '%s' in the filename as a " +
                        "placeholder for the beginning of the time interval")

    # parse
    options = ap.parse_args(args=args)

    if options.net is None:
        print("Error: a network to load must be given.")
        return 1
    if options.verbose:
        print("Reading network from '%s'" % options.net)
    net = sumolib.net.readNet(options.net, withInternal=options.internal)

    if options.measures is None:
        print("Error: a dump file must be given.")
        return 1

    dumps = options.dumps.split(",")
    times = []
    hc = None
    colorDump = dumps[0]
    colorMeasure = options.measures.split(",")[0]
    if colorDump:
        if options.verbose:
            print("Reading colors from '%s' (attribute:%s)" %
                    (colorDump, colorMeasure))
        hc = WeightsReader(colorMeasure)
        sumolib.output.parse_sax(colorDump, hc)
        times = hc._edge2value

    hw = None
    if len(dumps) > 1:
        widthDump = dumps[1]
        widthMeasure = options.measures.split(",")[1]
        if options.verbose:
            print("Reading width attribute from '%s' (attribute:%s)" %
                    (widthDump, widthMeasure))
        hw = WeightsReader(widthMeasure)
        sumolib.output.parse_sax(widthDump, hw)
        times = hw._edge2value

    # Should we also save the figure to a file / list of files (comma
    # separated)? Then we need to check the output filename(s)
    if options.output:
        options.nolegend = True
        optOutputNames = options.output

        # If we have multiple intervals to be plotted, make sure we have
        # proper output filenames (with a %s as a placeholder in it)
        if len(times) > 1 and optOutputNames.find('%s') < 0:
            print('Warning: multiple time intervals detected, but ' +
                    'the output filename(s) do not contain a \'%s\' placeholder. ' +
                    'Continuing by using a default placeholder.')

            # Modify each filename by putting a '-%s' right before the
            # extension
            filenames = optOutputNames.split(',')
            for i in range(0, len(filenames)):
                filename, extension = os.path.splitext(filenames[i])
                filenames[i] = filename + '-%s' + extension
            optOutputNames = ','.join(filenames)

    # Now go through each time interval and create the figures
    t = max(times)  # Draw last interval only
    # for t in times:
    if options.verbose:
        print("Processing interval with a beginning of %s" % t)
    colors = {}
    maxColorValue = None
    minColorValue = None
    for e in net._id2edge:
        if hc and t in hc._edge2value and e in hc._edge2value[t]:
            if options.colorMax is not None and hc._edge2value[t][e] > options.colorMax:
                hc._edge2value[t][e] = options.colorMax
            if options.colorMin is not None and hc._edge2value[t][e] < options.colorMin:
                hc._edge2value[t][e] = options.colorMin
            if maxColorValue is None or maxColorValue < hc._edge2value[t][e]:
                maxColorValue = hc._edge2value[t][e]
            if minColorValue is None or minColorValue > hc._edge2value[t][e]:
                minColorValue = hc._edge2value[t][e]
            colors[e] = hc._edge2value[t][e]
    if options.colorMax is not None:
        maxColorValue = options.colorMax
    if options.colorMin is not None:
        minColorValue = options.colorMin
    if options.logColors:
        helpers.logNormalise(colors, maxColorValue)
    else:
        helpers.linNormalise(colors, minColorValue, maxColorValue)
    for e in colors:
        colors[e] = helpers.getColor(options, colors[e], 1.)
    if options.verbose:
        print("Color values are between %s and %s" %
                (minColorValue, maxColorValue))

    widths = {}
    maxWidthValue = None
    minWidthValue = None
    for e in net._id2edge:
        if hw and t in hw._edge2value and e in hw._edge2value[t]:
            v = abs(hw._edge2value[t][e])
            if options.widthMax is not None and v > options.widthMax:
                v = options.widthMax
            if options.widthMin is not None and v < options.widthMin:
                v = options.widthMin
            if maxWidthValue is None or maxWidthValue < v:
                maxWidthValue = v
            if minWidthValue is None or minWidthValue > v:
                minWidthValue = v
            widths[e] = v
    if options.widthMax is not None:
        maxWidthValue = options.widthMax
    if options.widthMin is not None:
        minWidthValue = options.widthMin
    if options.logWidths:
        helpers.logNormalise(widths, options.colorMax)
    else:
        helpers.linNormalise(widths, minWidthValue, maxWidthValue)
    for e in widths:
        widths[e] = options.minWidth + widths[e] * \
            (options.maxWidth - options.minWidth)
    if options.verbose:
        print("Width values are between %s and %s" %
                (minWidthValue, maxWidthValue))
    if hw and (minWidthValue is None or maxWidthValue is None):
        print("Skipping interval %s without data" % t)

    fig, ax = helpers.openFigure(options)
    ax.set_aspect("equal", None, 'C')
    helpers.plotNet(net, colors, widths, options)

    # drawing the legend, at least for the colors
    norm = matplotlib.colors.LogNorm if options.logColors else matplotlib.colors.Normalize
    sm = plt.cm.ScalarMappable(cmap=matplotlib.cm.get_cmap(options.colormap),
                                norm=norm(vmin=minColorValue, vmax=maxColorValue))

    if sys.version_info.major < 3:
        # "fake up the array of the scalar mappable. Urgh..."
        # (pelson, http://stackoverflow.com/questions/8342549/matplotlib-add-colorbar-to-a-sequence-of-line-plots)
        sm._A = []
    color_bar = plt.colorbar(sm, ax=ax)
    color_bar.set_label(options.colorBarLabel)

    # Should we also save the figure to a file / list of files (comma
    # separated)?

    fig_path = f'{figfolder}heatmap_{int(t)}.png'
    fig.savefig(fig_path, dpi=1600)
    expandedOutputNames = None
    if options.output:

        # If we have a "%s" in the name of the output then replace it with the
        # interval begin of the current interval
        expandedOutputNames = optOutputNames
        if expandedOutputNames.find('%s') >= 0:
            expandedOutputNames = expandedOutputNames.replace("%s", str(t))

    # Can be used to print additional text in the figure:
    #
    # m, s = divmod(int(t), 60)
    # h, m = divmod(m, 60)
    # timeStr = "%02d:%02d:%02d" % (h, m, s)
    # ax.text(0.2, 0.2, timeStr, bbox={
    #     'facecolor': 'white', 'pad': 10}, size=16)
    helpers.closeFigure(fig, ax, options, False, expandedOutputNames)

    return fig_path

LLMAgent/plotHeatmap.py
lijlansg/TrafficGPT
from __future__ import absolute_import
from __future__ import print_function

import os
import sys

import matplotlib
matplotlib.use('Agg')

import sumolib  # noqa
from sumolib.visualization import helpers  # noqa
from sumolib.options import ArgumentParser  # noqa
import matplotlib.pyplot as plt  # noqa

def plot_intersections(target_junction_id, folderpath, args=None) -> str:
    """The main function; parses options and plots"""
    # ---------- build and read options ----------
    ap = ArgumentParser()
    ap.add_argument("-n", "--net", dest="net", category="input", type=ap.net_file, metavar="FILE",
                    required=True, help="Defines the network to read")
    ap.add_argument("-v", "--verbose", dest="verbose", action="store_true",
                    default=False, help="If set, the script says what it's doing")
    ap.add_argument("-w", "--width", dest="width",
                    type=float, default=20, help="Defines the width of the dots")
    ap.add_argument("--color", dest="color", category="visualization",
                    default='r', help="Defines the dot color")
    ap.add_argument("--edge-width", dest="defaultWidth", category="visualization",
                    type=float, default=1, help="Defines the edge width")
    ap.add_argument("--edge-color", dest="defaultColor", category="visualization",
                    default='k', help="Defines the edge color")
    # standard plot options
    helpers.addInteractionOptions(ap)
    helpers.addPlotOptions(ap)
    # parse
    options = ap.parse_args(args=args)

    if options.verbose:
        print("Reading network from '%s'" % options.net)
    net = sumolib.net.readNet(options.net)

    tlsn = {}
    for tid in net._id2tls:
        t = net._id2tls[tid]
        tlsn[tid] = set()
        for c in t._connections:
            n = c[0].getEdge().getToNode()
            tlsn[tid].add(n)

    tlspX = []
    tlspY = []
    junctionID = []
    for tid in tlsn:
        if tid in target_junction_id:
            x = 0
            y = 0
            n = 0
            for node in tlsn[tid]:
                x += node._coord[0]
                y += node._coord[1]
                n = n + 1
            x = x / n
            y = y / n
            tlspX.append(x)
            tlspY.append(y)
            junctionID.append(tid)

    fig, ax = helpers.openFigure(options)
    ax.set_aspect("equal", None, 'C')
    helpers.plotNet(net, {}, {}, options)
    plt.plot(tlspX, tlspY, options.color, linestyle='',
                marker='o', markersize=options.width)
    for i in range(len(junctionID)):
        plt.text(tlspX[i]*1.01, tlspY[i]*1.01, str(junctionID[i]), fontsize=10, color = "r", style = "italic", weight = "light", verticalalignment='center', horizontalalignment='right',rotation=0) #给散点加标签

    options.nolegend = True
    fig_path = f'{folderpath}intersections.png'
    fig.savefig(fig_path, dpi=1600)
    helpers.closeFigure(fig, ax, options)

    return fig_path

LLMAgent/plotIntersections.py
lijlansg/TrafficGPT
import pandas as pd
import xml.dom.minidom


def read_last_dump(dumpfile: str):
    dom = xml.dom.minidom.parse(dumpfile)
    root = dom.documentElement
    interval = root.getElementsByTagName('interval')[-1]
    edges = interval.getElementsByTagName('edge')
    LinkTDMOE = []
    for edge in edges:
        edge_id = edge.getAttribute('id')
        if eval(edge.getAttribute('sampledSeconds')) != 0:
            try:
                speed = eval(edge.getAttribute('speed'))
                waitingTime = eval(edge.getAttribute('waitingTime'))
                timeLoss = eval(edge.getAttribute('timeLoss'))
                left = eval(edge.getAttribute('left'))
                density = eval(edge.getAttribute('density'))
            except SyntaxError:
                speed = float('nan')
                waitingTime = float('nan')
                timeLoss = float('nan')
                left = float('nan')
                density = float('nan')
        else:
            speed = float('nan')
            waitingTime = float('nan')
            timeLoss = float('nan')
            left = float('nan')
            density = float('nan')
            # print(speed)
        LinkTDMOE.append(
            [edge_id, speed, waitingTime, timeLoss, left, density])

    df = pd.DataFrame(LinkTDMOE, columns=[
                        'edgeID', 'speed', 'waitingTime', 'timeLoss', 'left', 'density'])

    return df
LLMAgent/readDump.py
lijlansg/TrafficGPT

import os
import sys
import pandas as pd
import traci
from sumolib import checkBinary

from LLMAgent.buildGraph import Lane, Edge, Junction, Graph, build_graph
from LLMAgent.websterOptimize import Webster
from LLMAgent.plotIntersections import plot_intersections
from LLMAgent.plotHeatmap import plot_heatmap
from LLMAgent.readDump import read_last_dump


def prompts(name, description):
    def decorator(func):
        func.name = name
        func.description = description
        return func

    return decorator


class simulationControl:
    def __init__(self, sumocfgfile: str, netfile: str, dumpfile: str, originalstatefile: str, tempstatefile: str, figfolder: str) -> None:
        self.sumocfgfile = sumocfgfile
        self.netfile = netfile
        self.dumpfile = dumpfile
        self.originalstatefile = originalstatefile
        self.tempstatefile = tempstatefile
        self.figfolder = figfolder

    @prompts(name='Simulation Controller',
             description="""
             This tool is used to proceed and run the traffic simulation on SUMO. 
             The output will tell you whether you have finished this command successfully.
             This tool will also return the file path of a heat map of the road network as a supplementary information for you to provide the final answer. 
             The input should be a string, representing how many times have you called this tool, which shoule be a number greater than or equal to 0. 
             For example: if you never called this tool before and this is your first time calling this tool, the input should be 0; if you have called this tool twice, the imput should be 2.""")
    def inference(self, ordinal: str) -> str:
        ordinal_number = eval(ordinal)
        STEP = 600

        if 'SUMO_HOME' in os.environ:
            tools = os.path.join(os.environ['SUMO_HOME'], 'tools')
            sys.path.append(tools)
        else:
            raise RuntimeError(
                "please declare environment variable 'SUMO_HOME'")

        if_show_gui = False

        if not if_show_gui:
            sumoBinary = checkBinary('sumo')
        else:
            sumoBinary = checkBinary('sumo-gui')

        traci.start([sumoBinary, "-c", self.sumocfgfile])
        print('start reading state')
        if ordinal_number > 0:
            traci.simulation.loadState(self.tempstatefile)
        else:
            traci.simulation.loadState(self.originalstatefile)

        start_time = int(traci.simulation.getTime() / 1000)
        print('read state done!')
        for step in range(start_time, start_time + STEP):
            traci.simulationStep()

        traci.simulation.saveState(self.tempstatefile)

        traci.close()
        args = f'''-v -n {self.netfile} --measures speed,occupancy -i {self.dumpfile} \
            --default-width .5 --colormap RdYlGn  --max-width 3 --min-width .5 \
            --min-color-value 0 --max-color-value 15 --max-width-value 100 --min-width-value 0'''
        fig_path = plot_heatmap(self.figfolder, args)

        return f"You have successfully proceeded the traffic simulation on SUMO for 600 seconds. And your final answer should include this sentence without changing anything: the road network heat map is kept at: `{fig_path}`."


class intersectionPerformance:
    def __init__(self, netfile: str, dumpfile: str) -> None:
        self.netfile = netfile
        self.dumpfile = dumpfile

    @prompts(name='Get Intersection Performance',
             description="""
            This tool is used to get the traffic performance of all the intersections or several target intersections in the simulation road network.
            The output will provid traffic status information of intersetions in a tabular dataset.
            with various columns, such as Junction_id, speed_avg, volume_avg, and timeLoss_avg. Each row represents data for a specific junction. 
            Include the tabular dataset in markdown formart directly in your final answer! Do not try to change anything including the format! 
            The input If you do not have information of any specific intersection ID and the human didn't specify to get information of all the intersections, the input should be a string: 'None', and the tool will give you an overview data for the final answer so you don't need more specific information about certain intersections. 
            If you have specific target intersection IDs, the input should be a comma seperated string, with each part representing a target intersection ID.
            Only if you can find a word 'all' in the human message, the input can be a string: 'All'. """)
    def inference(self, target: str) -> str:

        if 'None' in target.replace(' ', '') or 'All' in target.replace(' ', ''):
            # print('no target' + target.replace(' ', ''))
            have_target = False
            target_junction_id = []
        else:
            have_target = True
            target_junction_id = target.replace(' ', '').split(',')
            # print('target'+ str(target.replace(' ', '').split(',')))

        graph = build_graph(self.netfile)
        edgedata = read_last_dump(self.dumpfile)

        junction_list = graph.junctions
        junction_summary_table = pd.DataFrame(
            columns=['Juction_id', 'speed_avg', 'volume_avg', 'timeLoss_avg'])
        for j_id, junction in junction_list.items():
            if len(junction.inEdges) <= 2:
                continue
            # print(j_id)
            if have_target and j_id not in target_junction_id:
                continue
            # print(j_id)
            upstream_list = []
            for edge in junction.inEdges:
                upstream_list.append(edge.id[0])
            junction_data = edgedata[edgedata["edgeID"].isin(upstream_list)]
            speed_avg = (
                junction_data['speed'] * junction_data['left']).sum() / junction_data['left'].sum()
            waitingTime_avg = (
                junction_data['waitingTime'] * junction_data['left']).sum() / junction_data['left'].sum()
            timeLoss_avg = (
                junction_data['timeLoss'] * junction_data['left']).sum() / junction_data['left'].sum()
            volume_avg = (
                junction_data['speed'] * 3.6 * junction_data['density']).mean()
            junction_summary_dic = {"Juction_id": j_id, "speed_avg": speed_avg,
                                    "volume_avg": volume_avg, "timeLoss_avg": timeLoss_avg}
            new_row = pd.DataFrame(junction_summary_dic, index=[0])
            junction_summary_table = pd.concat(
                [junction_summary_table, new_row], axis=0).reset_index(drop=True)
            # print(junction_summary_dic)
        sorted_table = junction_summary_table.sort_values(
            by=['speed_avg', 'volume_avg', 'timeLoss_avg'], ascending=[True, False, False]).reset_index(drop=True)
        # print(sorted_table)
        if 'None' in target.replace(' ', ''):
            msg = 'No specific target intersections. So, I can show you the overview by providing the traffic status of 5 intersections in the worst operating condition by default. Make sure you output the tabular content in markdown format into your final answer. \n'
            return msg + sorted_table.head().to_markdown()
        elif 'All' in target.replace(' ', ''):
            msg = 'Here are the traffic status of all intersections. Make sure you output the tabular content in markdown format into your final answer. \n'
            return msg + sorted_table.to_markdown()
        else:
            msg = 'Here are the traffic status of your targeted intersections. Make sure you output the tabular content in markdown format into your final answer. \n'
            return msg + sorted_table.to_markdown()


class intersectionSignalOptimization:
    def __init__(self, netfile: str, configfile: str, routefile: str, tlsfile: str) -> None:
        self.netfile = netfile
        self.configfile = configfile
        self.routefile = routefile
        self.tlsfile = tlsfile

    @prompts(name='Optimize Intersection Signal Control Scheme',
             description="""
            This tool is used to optimize the signal control scheme of several target intersections in the simulation road network.
            Do not use this tool unless the human user asks to optimize intersections.
            The output will tell you whether you have finished this command successfully.
            The input should be a comma seperated string, with each part representing a target intersection ID. """)
    def inference(self, target: str) -> str:

        if 'None' in target:
            return "Please provide the target intersection IDs."

        options = f'-n {self.netfile} -f {self.configfile} -r {self.routefile} -o {self.tlsfile}'
        target_ID = target.replace(' ', '').split(',')

        optimizer = Webster(target_ID, options)
        optimizer.run_webster()
        optimizer.add_TLS()

        return f"The signal control scheme for intersection {target} has already been optimized successfully according to Webster's algorithm. The new TLS program has been written to the configuration file. If you want to see the traffic status after the optimization, you need to run the simulation again."


class intersectionVisulization:
    def __init__(self, netfile: str, figfolder: str) -> None:
        self.netfile = netfile
        self.figfolder = figfolder

    @prompts(name='Visualize Intersections',
             description="""
            This tool is used to show the locations of several target intersections by visualize them on a map.
            Use this tool more than others if the question is about locations of intersections.
            The output will tell you whether you have finished this command successfully.
            The input should be a comma seperated string, with each part representing a target intersection ID. """)
    def inference(self, target: str) -> str:

        target_junction_id = target.replace(' ', '').split(',')
        options = f'-n {self.netfile} --width 5 --edge-color #606060'

        fig_path = plot_intersections(
            target_junction_id,
            self.figfolder,
            options
        )

        return f"You have successfully visualized the location of intersection {target} on the following map. And your final answer should include this sentence without changing anything: The location of intersection {target} is kept at: `{fig_path}`."

LLMAgent/trafficTools.py
lijlansg/TrafficGPT
"""
- The Webster's equation is used to optimize the cycle length
  and the green times of the traffic lights at common intersections
  in a sumo network with a given route file (vehicles and bicycles).

- One hour flow is considered. If the flow arrival situation within
  1 hour is "turbulent", i.e. 1-h flow arrives within 15 min, then
  a poor result is expected.

- The signal at each intersection is optimized independently. Thus,
  some poor result may occur if there is conflict between the singal
  programs of different intersections. Signal coordination is not
  handled here, but in tlsCoordinator.py.

- The parameters in the Webster's equation may need to be adjusted
  according to the respective user case. Otherwise, worse results
  may occur. E.g. minimal green time, minimal cycle length and
  saturation headway. The latter one can affect the calcuated
  road capacity.

- The allowed movements per lane in the network should correspond to
  those in the respective signal group/phase. E.g. if through and
  right movements are allowed in a phase, they should be able to
  share their lanes without conflicts.  Otherwise, worse results
  may occur.

- Traffic lights without traffic flows will not be optimized.

- Flows are considered only for one hour, corresponding to
  the Webster's equation. Using the option --begin to define the
  start time of the investigated one hour.

- PCE is used instead of the number of vehicles.

- If a lane group has green times in more than one phase,
  the respective flows will be equally divided into the corresponding
  phases for calculating the green splits.

- If multiple major greens exist, the respective flows will be evenly
  distributed in each major green.

- Flows only with minor greens at a certain direction is considered.

- If the critical flow or the sum of the critical flows is larger than 1,
  the optimal cycle length will be set to 120 sec which can be adjusted
  in the option.

- Duration for yellow phase will be adjusted according to
  the defined option value

- Duration for all-red phase will not be adjusted.

- Pedestrians are not considered yet.

- Triggered vehicles are not considered.

- The begin time is the first vehicle's departure time, if it is not given.

- If the period is less than 1 hour, the flow will be scaled up to 1-hour flow.

- If the period is longer than 1 hour and no begin time is given, the 1-hour
  peak-flow will be used.
"""

from __future__ import absolute_import
from __future__ import print_function

import collections
import sys
import os
import sumolib



def getEdges(veh, net, routes):
    if isinstance(veh.route, str):
        edges = routes[veh.route]
    else:
        edges = veh.route[0].edges
    return [net.getEdge(e) for e in edges.split()]

def getRoutes(files):
    route_dict = {}
    for file in files:
        for route in sumolib.output.parse(file, 'route'):
            route_dict[route.id] = route.edges
    return route_dict

def getFlows( net, routeFiles, tlsList, begin, scale_fac, verbose, isSorted=False):
    tlsFlowsMap = {}
    end = begin + 3600
    for tls in tlsList:
        tlsFlowsMap[tls._id] = collections.defaultdict(
            lambda: collections.defaultdict(int))
    route_dict = getRoutes(routeFiles.split(','))
    for file in routeFiles.split(','):
        if verbose:
            print("parsing route file:", file)
        triggered = parsed = 0
        for veh in sumolib.output.parse(file, 'vehicle'):
            if veh.depart == "triggered":
                triggered += 1
                continue
            if sumolib.miscutils.parseTime(veh.depart) >= end:
                if isSorted:
                    break
                continue
            if sumolib.miscutils.parseTime(veh.depart) >= begin:
                edgeList = getEdges(veh, net, route_dict)
                for idx, edge in enumerate(edgeList):
                    tls = None if edge.getToNode().getType() in (
                        "rail_crossing", "rail_signal") else edge.getTLS()
                    if tls and idx < len(edgeList) - 1:
                        # c: [[inLane, outLane, linkNo],[],..]
                        for c in tls.getConnections():
                            inEdge = c[0].getEdge()
                            outEdge = c[1].getEdge()
                            if inEdge == edge and outEdge == edgeList[idx + 1]:
                                pce = 1.
                                if veh.type == "bicycle":
                                    pce = 0.2
                                elif veh.type in ["moped", "motorcycle"]:
                                    pce = 0.5
                                elif veh.type in ["truck", "trailer", "bus", "coach"]:
                                    pce = 3.5
                                tlsFlowsMap[tls.getID()][inEdge.getID(
                                ) + " " + outEdge.getID()][c[2]] += pce
                                parsed += 1
        if triggered > 0:
            print("Warning: Ignored %s triggered vehicles in %s." %
                    (triggered, file))
        if parsed == 0:
            print("Warning: No vehicles parsed from %s." % file)
        elif verbose:
            print("Parsed %s vehicles from %s." % (parsed, file))
    # scale up flow
    if scale_fac != 1.:
        for t in tlsList:
            for subRoute in tlsFlowsMap[t.getID()]:
                for conn in tlsFlowsMap[t.getID()][subRoute]:
                    tlsFlowsMap[t.getID()][subRoute][conn] *= scale_fac

    # remove the doubled counts
    connFlowsMap = {}
    for t in tlsList:
        connFlowsMap[t.getID()] = {}
        for subRoute in tlsFlowsMap[t.getID()]:
            totalConns = len(tlsFlowsMap[t.getID()][subRoute])
            for conn in tlsFlowsMap[t.getID()][subRoute]:
                tlsFlowsMap[t.getID()][subRoute][conn] /= totalConns
                connFlowsMap[t.getID()][conn] = tlsFlowsMap[t.getID()
                                                            ][subRoute][conn]
        # remove the redundant connection flows
        connFlowsMap = removeRedundantFlows(t, connFlowsMap)

    return connFlowsMap

def getEffectiveTlsList(tlsList, connFlowsMap, skipList, verbose, TARGET_ID):
    effectiveTlsList = []
    tlsList = [tl for tl in tlsList if tl._id not in skipList]
    for tl in tlsList:
        if tl._id not in TARGET_ID:
            continue
        if len(tl.getPrograms()) == 0:
            continue

        for program in tl.getPrograms().values():
            for phase in program.getPhases():
                if len(phase.state) > len(tl.getConnections()):
                    print("Warning: the number of unused states at TLS %s: %s (%s states, %s connections)" %
                            (tl.getID(), (len(phase.state)-len(tl.getConnections())),
                            len(phase.state), len(tl.getConnections())))
                    break

        for conn in connFlowsMap[tl.getID()]:
            if connFlowsMap[tl.getID()][conn] > 0:
                effectiveTlsList.append(tl)
                break
    return effectiveTlsList

def removeRedundantFlows(t, connFlowsMap):
    # if two or more intersections share the lane-lane connection indices together,
    # the redundant connection flows will set to zero.
    connsList = t.getConnections()
    connsList = sorted(connsList, key=lambda connsList: connsList[2])
    redundantConnsList = []
    identical = True
    for c1 in connsList:
        for c2 in connsList:
            if c1[2] != c2[2]:
                if c1[1]._edge == c2[0]._edge:
                    identical = identityCheck(
                        c1[0]._edge, c2[0]._edge._incoming, identical)
                    if identical:
                        for toEdge in c2[0]._edge._outgoing:
                            for c in c2[0]._edge._outgoing[toEdge]:
                                if c._tlLink not in redundantConnsList:
                                    redundantConnsList.append(c._tlLink)
                    else:
                        for conn_1 in c1[0]._edge._outgoing[c2[0]._edge]:
                            if conn_1._direction == 's':
                                for toEdge in c2[0]._edge._outgoing:
                                    for conn_2 in c2[0]._edge._outgoing[toEdge]:
                                        if conn_2._tlLink not in redundantConnsList:
                                            redundantConnsList.append(
                                                conn_2._tlLink)
    for conn in redundantConnsList:
        if conn in connFlowsMap[t._id]:
            connFlowsMap[t._id][conn] = 0.
    return connFlowsMap

def identityCheck(e1, incomingLinks, identical):
    for i in incomingLinks:
        if i != e1:
            identical = False
            break
    return identical

def getLaneGroupFlows(tl, connFlowsMap, phases, greenFilter, multiOwnGreenMap, getmultiOwnGreen, options):
    connsList = tl.getConnections()
    groupFlowsMap = {}  # i(phase): duration, laneGroup1, laneGroup2, ...
    # tls-linkIndex: connsList[i][2]
    connsList = sorted(connsList, key=lambda connsList: connsList[2])
    aktiveLinkIndices = set()

    # get acktive link indices that used in the current TLS plan
    # the connections with internal links for walking areas and crossings are also considered
    for conn in connsList:
        for toEdge in conn[0]._edge._outgoing:
            for c in conn[0]._edge._outgoing[toEdge]:
                if c._tlLink >= 0 and toEdge._function not in ["crossing", "walkingarea", "internal"]:
                    aktiveLinkIndices.add(c._tlLink)

    # check if there are shared lane groups, i.e. some lane groups have only "g" (no "G")
    ownGreenConnsList = []   # connections with major green
    for i, p in enumerate(phases):
        # j is the linkIndex
        for j, control in enumerate(p.state):
            if control == "G" and j in aktiveLinkIndices:
                if j not in ownGreenConnsList:
                    ownGreenConnsList.append(j)
                elif not getmultiOwnGreen:       # j could be put more than once in the map
                    multiOwnGreenMap[tl._id].append(j)
    if options.verbose:
        if multiOwnGreenMap[tl._id]:
            print('TLS: %s --> the tl-indices with more than one major-green:%s' %
                    (tl._id, multiOwnGreenMap[tl._id]))
    yellowRedTime = 0
    greenTime = 0
    currentLength = 0
    phaseLaneIndexMap = collections.defaultdict(list)

    for i, p in enumerate(phases):
        currentLength += p.duration
        if 'G' in p.state and 'y' not in p.state and p.duration >= greenFilter:
            greenTime += p.duration
            groupFlowsMap[i] = [p.duration]
            groupFlows = 0
            laneIndexList = []
            exEdge = None
            for j, control in enumerate(p.state):
                if j in aktiveLinkIndices:
                    inEdge = connsList[j][0]._edge._id
                    multiOwnGreenFactor = 1
                    multiOwnGreenFactor += multiOwnGreenMap[tl._id].count(
                        j)
                    processed = False
                    if not exEdge:
                        exEdge = inEdge
                    # protected green directly after major green for the same edge
                    if ((inEdge == exEdge) and ((control == 'G') or (control == 'g' and j not in ownGreenConnsList))):
                        if j in connFlowsMap[tl._id]:
                            # if a connection flow has more than one major green,
                            # the flow is regularly distributed in each "major green"
                            groupFlows += connFlowsMap[tl._id][j] / \
                                float(multiOwnGreenFactor)

                        if connsList[j][0].getIndex() not in laneIndexList:
                            laneIndexList.append(
                                connsList[j][0].getIndex())
                        processed = True

                    # fromEdge is different from the previous one or the last state
                    if exEdge != inEdge or j == len(p.state) - 1:
                        # save the data of the previous group
                        if laneIndexList:
                            phaseLaneIndexMap[i].append(laneIndexList)
                            groupFlowsMap[i].append(groupFlows)
                            # reset
                            laneIndexList = []
                            groupFlows = 0
                            if (j == len(p.state) - 1) and processed:
                                break

                        if control == "G":
                            if connsList[j][0].getIndex() not in laneIndexList:
                                laneIndexList.append(
                                    connsList[j][0].getIndex())
                            if j in connFlowsMap[tl._id]:  # only flows > 0
                                groupFlows += connFlowsMap[tl._id][j] / \
                                    float(multiOwnGreenFactor)
                        elif control == 'g' and j not in ownGreenConnsList:
                            if connsList[j][0].getIndex() not in laneIndexList:
                                laneIndexList.append(
                                    connsList[j][0].getIndex())
                            if j in connFlowsMap[tl._id]:  # only flows > 0
                                groupFlows += connFlowsMap[tl._id][j] / \
                                    float(multiOwnGreenFactor)

                        if (j == len(p.state) - 1) and laneIndexList:
                            phaseLaneIndexMap[i].append(laneIndexList)
                            groupFlowsMap[i].append(groupFlows)

                            # reset
                            laneIndexList = []
                            groupFlows = 0

                    exEdge = inEdge
        elif 'G' not in p.state and 'g' in p.state and 'y' not in p.state and 'r' not in p.state:
            print("Check: only g for all connections:%s in phase %s" %
                    (tl._id, i))
        elif ('G' not in p.state and 'g' not in p.state) or ('G' not in p.state and 'y' in p.state and 'r' in p.state):
            yellowRedTime += int(p.duration)
        if options.verbose and i in groupFlowsMap:
            print("phase: %s" % i)
            print("group flows: %s" % groupFlowsMap[i])
            print("The used lanes: %s" % phaseLaneIndexMap[i])
    if options.verbose:
        print("the current cycle length:%s sec" % currentLength)
    return groupFlowsMap, phaseLaneIndexMap, currentLength, multiOwnGreenMap, getmultiOwnGreen

def getMaxOptimizedCycle(groupFlowsMap, phaseLaneIndexMap, currentLength, cycleList, multiOwnGreenMap, options):
    lostTime = len(groupFlowsMap) * options.losttime + options.allred
    satFlows = 3600. / options.satheadway
    # calculate the critical flow ratios and the respective sum
    criticalFlowRateMap = {}
    for i in groupFlowsMap:   # [duration. groupFlow1, groupFlow2...]
        criticalFlowRateMap[i] = 0.
        maxFlow = 0
        index = None
        if len(groupFlowsMap[i][1:]) > 0:
            for j, f in enumerate(groupFlowsMap[i][1:]):
                if f >= maxFlow:
                    maxFlow = f
                    index = j
            criticalFlowRateMap[i] = (
                maxFlow / float((len(phaseLaneIndexMap[i][index])))) / satFlows
        else:
            criticalFlowRateMap[i] = 0.
    sumCriticalFlows = sum(criticalFlowRateMap.values())

    if options.existcycle:
        optCycle = currentLength
    elif sumCriticalFlows >= 1.:
        optCycle = options.maxcycle
        if options.verbose:
            print("Warning: the sum of the critical flows >= 1:%s" %
                    sumCriticalFlows)
    else:
        optCycle = int(
            round((1.5 * lostTime + 5.) / (1. - sumCriticalFlows)))

    if not options.existcycle and optCycle < options.mincycle:
        optCycle = options.mincycle
    elif not options.existcycle and optCycle > options.maxcycle:
        optCycle = options.maxcycle

    cycleList.append(optCycle)

    return cycleList

def optimizeGreenTime(tl, groupFlowsMap, phaseLaneIndexMap, currentLength, multiOwnGreenMap, options):
    lostTime = len(groupFlowsMap) * options.losttime + options.allred
    satFlows = 3600. / options.satheadway
    # calculate the critical flow ratios and the respective sum
    criticalFlowRateMap = {}
    for i in groupFlowsMap:   # [duration. groupFlow1, groupFlow2...]
        criticalFlowRateMap[i] = 0.
        maxFlow = 0
        index = None
        if len(groupFlowsMap[i][1:]) > 0:
            for j, f in enumerate(groupFlowsMap[i][1:]):
                if f >= maxFlow:
                    maxFlow = f
                    index = j
            criticalFlowRateMap[i] = (
                maxFlow / float((len(phaseLaneIndexMap[i][index])))) / satFlows
        else:
            criticalFlowRateMap[i] = 0.
    sumCriticalFlows = sum(criticalFlowRateMap.values())
    if options.write_critical_flows:
        print(tl.getID(), criticalFlowRateMap)
        print('sum of the critical flow ratios: ', sumCriticalFlows)

    if options.existcycle:
        optCycle = currentLength
    elif sumCriticalFlows >= 1.:
        optCycle = options.maxcycle
        if options.verbose:
            print("Warning: the sum of the critical flows >= 1:%s" %
                    sumCriticalFlows)
            print(
                'Warning: the maximal cycle defined in the option will be used as the optimal cycle.')
    else:
        optCycle = int(
            round((1.5 * lostTime + 5.) / (1. - sumCriticalFlows)))

    if not options.existcycle and optCycle < options.mincycle:
        optCycle = options.mincycle
    elif not options.existcycle and optCycle > options.maxcycle:
        optCycle = options.maxcycle

    # calculate the green time for each critical group
    effGreenTime = optCycle - lostTime
    totalLength = len(groupFlowsMap)*options.yellowtime + options.allred
    minGreenPhasesList = []
    adjustGreenTimes = 0
    totalGreenTimes = 0
    subtotalGreenTimes = 0

    for i in criticalFlowRateMap:
        groupFlowsMap[i][0] = effGreenTime * \
            (criticalFlowRateMap[i] / sum(criticalFlowRateMap.values())
                ) - options.yellowtime + options.losttime
        groupFlowsMap[i][0] = int(round(groupFlowsMap[i][0]))
        totalGreenTimes += groupFlowsMap[i][0]
        if groupFlowsMap[i][0] < options.mingreen:
            groupFlowsMap[i][0] = options.mingreen
            minGreenPhasesList.append(i)
        else:
            subtotalGreenTimes += groupFlowsMap[i][0]
        totalLength += groupFlowsMap[i][0]

    # adjust the green times if minimal green times are applied for keeping the defined maximal cycle length.
    if minGreenPhasesList and totalLength > options.maxcycle and options.restrict:
        totalLength = len(groupFlowsMap) * \
            options.yellowtime + options.allred
        if options.verbose:
            print("Re-allocate the green splits!")
        adjustGreenTimes = totalGreenTimes - \
            len(minGreenPhasesList) * options.mingreen
        for i in groupFlowsMap:
            if i not in minGreenPhasesList:
                groupFlowsMap[i][0] = int(
                    round((groupFlowsMap[i][0] / float(subtotalGreenTimes)) * adjustGreenTimes))
            totalLength += groupFlowsMap[i][0]

    if options.unified_cycle and totalLength != optCycle:
        diff = optCycle - totalLength
        secs_to_distribute = [int(diff / abs(diff))] * abs(diff)
        keys = list(groupFlowsMap.keys())
        for i, s in enumerate(secs_to_distribute):
            groupFlowsMap[keys[i % len(groupFlowsMap)]][0] += s

    if options.verbose:
        totalLength = len(groupFlowsMap) * \
            options.yellowtime + options.allred
        for i in groupFlowsMap:
            totalLength += groupFlowsMap[i][0]
            print("Green time for phase %s: %s" % (i, groupFlowsMap[i][0]))
        print("The optimal cycle length:%s\n" % totalLength)

    return groupFlowsMap

def checkRoutePeriod(routefiles, begin):
    scale_fac = 1.
    # find the last vehicle's departure time
    veh_starttime = -1.
    veh_endtime = -1.
    for file in routefiles.split(','):
        for veh in sumolib.output.parse(file, 'vehicle'):
            if veh.depart != "triggered":
                depart = float(veh.depart)
                if veh_starttime == -1.:
                    veh_starttime = veh_endtime = depart

                if veh_endtime < depart:
                    veh_endtime = depart
                if veh_starttime > depart:
                    veh_starttime = depart

    # check the begin time
    checkPeak = False
    if begin is None or begin < 0.:
        print("Warning: The begin time '%s' is not valid." % begin)
        begin = veh_starttime
        print(
            "Warning: The begin time is set to the first vehicle's departure time: %s." % veh_starttime)
        checkPeak = True

    # check how to process the flow
    end = begin + 3599.
    if veh_endtime < end:
        scale_fac += (end - veh_endtime)/3600.
        print("Warning: The period is less than 1 hour. "
                "The flows will be proportionally scaled up to 1-hour flow with the scaling factor %s." % scale_fac)
    elif checkPeak and veh_endtime > end:
        begin, peakFlow = getPeakFlowBegin(
            routefiles, begin, veh_endtime)
        print("Warning: The period (begining with %s) with the peak flow (%s) is used." % (
            begin, peakFlow))

    return begin, scale_fac

def getPeakFlowBegin(routefiles, begin, veh_endtime):
    minuteFlowMap = {}
    end_intl = int(veh_endtime//60.)+1
    for n in range(0, end_intl+1):
        minuteFlowMap[n] = 0.
    peak_begin = begin
    peakFlow = 0.

    for file in routefiles.split(','):
        for veh in sumolib.output.parse(file, 'vehicle'):
            if veh.depart != "triggered" and sumolib.miscutils.parseTime(veh.depart) >= begin:
                pce = 1.
                if veh.type == "bicycle":
                    pce = 0.2
                elif veh.type in ["moped", "motorcycle"]:
                    pce = 0.5
                elif veh.type in ["truck", "trailer", "bus", "coach"]:
                    pce = 3.5
                intl = int(float(veh.depart)//60.)
                minuteFlowMap[intl] += pce

    for i in range(0, end_intl-59):
        temp_sum = 0
        sub_end = i + 60
        if sub_end > end_intl:
            print(
                "Warning: The end time is larger and set to the last vehicle's departure time(%s)" % end_intl)
            sub_end = end_intl+1
        for j in range(i, sub_end):
            temp_sum += minuteFlowMap[j]
        if temp_sum > peakFlow:
            peakFlow = temp_sum
            peak_begin = i

    return peak_begin, peakFlow

class Webster:
    """韦伯斯特信号配时优化"""
 
    def __init__(self, target_ID, args):
        """初始化options"""
        self.target_ID = target_ID
        self.options = self.get_options(args)

    def get_options(self, args=None):
        optParser = sumolib.options.ArgumentParser()
        optParser.add_option("-n", "--net-file", category="input", dest="netfile", required=True,
                             help="define the net file (mandatory)")
        optParser.add_option("-f", "--config-file", category="input", dest="configfile", required=True,
                             help="define the config file (mandatory)")
        optParser.add_option("-o", "--output-file", category="output", dest="outfile",
                             default="tlsAdaptation.add.xml", help="define the output filename")
        optParser.add_option("-r", "--route-files", category="input", dest="routefiles", required=True,
                             help="define the route file separated by comma (mandatory)")
        optParser.add_option("-b", "--begin", category="time", dest="begin", type=float,
                             help="begin time of the optimization period with unit second (mandatory)")
        optParser.add_option("-y", "--yellow-time", category="processing", dest="yellowtime", type=int,
                             default=4, help="yellow time")
        optParser.add_option("-a", "--all-red", category="processing", dest="allred", type=int,
                             default=0, help="all-red time per cycle")
        optParser.add_option("-l", "--lost-time", category="processing", dest="losttime", type=int,
                             default=4, help="lost time for start-up and clearance in each phase")
        optParser.add_option("-g", "--min-green", category="processing", dest="mingreen", type=int,
                             default=4, help=" minimal green time when there is no traffic volume")
        optParser.add_option("--green-filter-time", category="processing", dest="greenFilter", type=int, default=0,
                             help="when computing critical flows, do not count phases with a green time below INT")
        optParser.add_option("--min-cycle", category="processing", dest="mincycle", type=int,
                             default=20, help="minimal cycle length")
        optParser.add_option("--max-cycle", category="processing", dest="maxcycle", type=int,
                             default=120, help="maximal cycle length")
        optParser.add_option("-e", "--existing-cycle", category="processing", dest="existcycle", action="store_true",
                             default=False, help="use the existing cycle length")
        optParser.add_option("--write-critical-flows", category="processing", dest="write_critical_flows",
                             action="store_true", default=False, help="print critical flows for each tls and phase")
        optParser.add_option("-p", "--program", category="processing", dest="program", default="a",
                             help="save new definitions with this program id")
        optParser.add_option("-H", "--saturation-headway", category="processing", dest="satheadway", type=float, default=2,
                             help="saturation headway in seconds for calculating hourly saturation flows")
        optParser.add_option("-R", "--restrict-cyclelength", category="processing", dest="restrict", action="store_true",
                             default=False, help="restrict the max. cycle length as the given one")
        optParser.add_option("-u", "--unified-cycle", category="processing", action="store_true", default=False,
                             help="use the calculated max cycle length as the cycle length for all intersections")
        optParser.add_option("--sorted", category="processing", action="store_true", default=False,
                             help="assume the route file is sorted (aborts reading earlier)")
        optParser.add_option("--skip", category="processing", dest="skip",
                             default='', help="the tls ids, which are skipped and seperated by comma")
        optParser.add_option("-v", "--verbose", category="processing", dest="verbose", action="store_true",
                             default=False, help="tell me what you are doing")
        return optParser.parse_args(args=args)

    def run_webster(self):
        if not self.options.netfile or not self.options.routefiles:
            raise RuntimeError("Error: Either both the net file and the route file or one of them are/is missing.")

        # check the period of the given route files, and find the peak-flow period if necessary
        begin, scale_fac = checkRoutePeriod(
            self.options.routefiles, self.options.begin)

        net = sumolib.net.readNet(
            self.options.netfile, withPrograms=True, withPedestrianConnections=True)
        tlsList = net.getTrafficLights()
        skipList = []
        skipList = self.options.skip.split(",")

        if self.options.verbose:
            print("the total number of tls: %s" % len(tlsList))

        print("Begin time:%s" % begin)
        # get traffic flows for each connection at each TL
        connFlowsMap = getFlows(
            net, self.options.routefiles, tlsList, begin, scale_fac, self.options.verbose, self.options.sorted)

        # remove the tls where no traffic volumes exist and which should be skipped
        effectiveTlsList = getEffectiveTlsList(
            tlsList, connFlowsMap, skipList, self.options.verbose, self.target_ID)
        print(len(effectiveTlsList))

        # prepare a map for any connection flow has more than one major-green in states
        multiOwnGreenMap = {}
        getmultiOwnGreen = False
        for tl in effectiveTlsList:
            multiOwnGreenMap[tl._id] = []

        with open(self.options.outfile, 'w') as outf:
            sumolib.xml.writeHeader(outf, root="additional", options=self.options)
            if len(effectiveTlsList) > 0:
                if self.options.unified_cycle:
                    cycleList = []
                    if self.options.verbose:
                        print(
                            "Firstly only calculate the maximal optimized cycle length! ")
                    for tl in effectiveTlsList:
                        if self.options.verbose:
                            print("tl-logic ID: %s" % tl._id)
                        programs = tl.getPrograms()
                        for pro in programs:
                            phases = programs[pro].getPhases()

                            # get the connection flows and group flows
                            r = getLaneGroupFlows(
                                tl, connFlowsMap, phases, 0, multiOwnGreenMap, getmultiOwnGreen, self.options)
                            groupFlowsMap, phaseLaneIndexMap, currentLength, multiOwnGreenMap, getmultiOwnGreen = r

                            # only optimize the cycle length
                            cycleList = getMaxOptimizedCycle(groupFlowsMap, phaseLaneIndexMap,
                                                                    currentLength, cycleList, multiOwnGreenMap, self.options)

                    self.options.maxcycle = max(cycleList)
                    self.options.mincycle = max(cycleList)
                    self.options.restrict = True
                    if self.options.verbose:
                        print("The maximal optimized cycle length is %s." %
                                max(cycleList))
                        print(
                            " It will be used for calculating the green splits for all intersections.")

                # calculate the green splits;
                # the optimal length will be also calculate if options.unified_cycle is set as false.
                for tl in effectiveTlsList:
                    if self.options.verbose:
                        print("tl-logic ID: %s" % tl._id)
                    programs = tl.getPrograms()
                    for pro in programs:
                        phases = programs[pro].getPhases()

                        # get the connection flows and group flows
                        r = getLaneGroupFlows(tl, connFlowsMap, phases, self.options.greenFilter,
                                                    multiOwnGreenMap, getmultiOwnGreen, self.options)
                        groupFlowsMap, phaseLaneIndexMap, currentLength, multiOwnGreenMap, getmultiOwnGreen = r

                        # optimize the cycle length and calculate the respective green splits
                        groupFlowsMap = optimizeGreenTime(tl, groupFlowsMap, phaseLaneIndexMap,
                                                                currentLength, multiOwnGreenMap, self.options)

                    # write output
                    outf.write('    <tlLogic id="%s" type="%s" programID="%s" offset="%i">\n' %
                                (tl._id, programs[pro]._type, self.options.program, programs[pro]._offset))

                    phases = programs[pro].getPhases()
                    for i, p in enumerate(phases):
                        duration = p.duration
                        if i in groupFlowsMap:
                            duration = groupFlowsMap[i][0]
                        # the yellow phase
                        elif 'y' in p.state and 'r' in p.state:
                            duration = self.options.yellowtime
                        else:
                            print(
                                "Duration for Phase %s is from the input file." % i)
                        outf.write(
                            '        <phase duration="%s" state="%s"/>\n' % (duration, p.state))
                    outf.write('    </tlLogic>\n')
            else:
                print(
                    "There are no flows at the given intersections. No green time optimization is done.")
            outf.write('</additional>\n')

    def add_TLS(self):
        tls_file = os.path.basename(self.options.outfile)
        oldStr1 = '<additional-files value="add.add.xml"/>'
        newStr1 = f'<additional-files value="add.add.xml,{tls_file}"/>'

        oldStr2 = f'<additional-files value="add.add.xml,{tls_file}"/>'

        with open(self.options.configfile, "r", encoding="utf-8") as f:
            lines = f.readlines()  # 将文件内容保存到内存
        with open(self.options.configfile, "w", encoding="utf-8") as f_w:
            # print(oldStr)
            for line in lines:
                if oldStr2 in line:  # 如果已经添加了tls_file，不用改变config文件的内容，因为新的信号灯方案只需在tls_file中更新 --2023.8.2 zsy
                    f_w.write(line)
                    continue
                if oldStr1 in line:  # 如果没有tls_file，将其添加到config文件中
                    line = line.replace(oldStr1, newStr1)
                f_w.write(line)
        print(
            f"The new TLS program has writen in the config file: {self.options.configfile}.")

LLMAgent/websterOptimize.py
lijlansg/TrafficGPT
import os
import re
import yaml
from rich import print
# from langchain import OpenAI
from langchain.chat_models import AzureChatOpenAI, ChatOpenAI

from LLMAgent.ConversationBot import ConversationBot

from LLMAgent.trafficTools import (
    simulationControl,
    intersectionPerformance,
    intersectionSignalOptimization,
    intersectionVisulization,
)

import gradio as gr
import openai.api_requestor
openai.api_requestor.TIMEOUT_SECS = 30

# ------------------------------------------------------------------------------
# --ZH 初始化 LLM
# --EN Initialize a LLM
OPENAI_CONFIG = yaml.load(open('config.yaml'), Loader=yaml.FullLoader)
if OPENAI_CONFIG['OPENAI_API_TYPE'] == 'azure':
    os.environ["OPENAI_API_TYPE"] = OPENAI_CONFIG['OPENAI_API_TYPE']
    os.environ["OPENAI_API_VERSION"] = OPENAI_CONFIG['AZURE_API_VERSION']
    os.environ["OPENAI_API_BASE"] = OPENAI_CONFIG['AZURE_API_BASE']
    os.environ["OPENAI_API_KEY"] = OPENAI_CONFIG['AZURE_API_KEY']
    llm = AzureChatOpenAI(
        deployment_name=OPENAI_CONFIG['AZURE_MODEL'],
        temperature=0,
        max_tokens=1024,
        request_timeout=60
    )
elif OPENAI_CONFIG['OPENAI_API_TYPE'] == 'openai':
    os.environ["OPENAI_API_KEY"] = OPENAI_CONFIG['OPENAI_KEY']
    llm = ChatOpenAI(
        temperature=0,
        model_name='gpt-3.5-turbo-16k-0613',  # or any other model with 8k+ context
        max_tokens=1024,
        request_timeout=60
    )

# ------------------------------------------------------------------------------
# --ZH 初始化工具
# --EN Initialize the tools

if not os.path.exists('./fig/'):
    os.mkdir('./fig/')


sumoCFGFile = './real-world-simulation-withTLS/xuancheng.sumocfg'
sumoNetFile = './real-world-simulation-withTLS/xuancheng.net.xml'
sumoRouFile = './real-world-simulation-withTLS/xuancheng.net.xml'
sumoEdgeDataFile = './real-world-simulation-withTLS/edgedata.xml'
sumoOriginalStateFile = './real-world-simulation-withTLS/originalstate.xml'
sumoTempStateFile = './real-world-simulation-withTLS/tempstate.xml'
sumoNewTLSFile = './real-world-simulation-withTLS/newTLS.add.xml'
targetFilePath = './fig/'
toolModels = [
    simulationControl(
        sumoCFGFile, sumoNetFile, sumoEdgeDataFile,
        sumoOriginalStateFile, sumoTempStateFile, targetFilePath
    ),
    intersectionPerformance(sumoNetFile, sumoEdgeDataFile),
    intersectionSignalOptimization(
        sumoNetFile, sumoCFGFile, sumoRouFile, sumoNewTLSFile,
    ),
    intersectionVisulization(sumoNetFile, targetFilePath)
]

# ------------------------------------------------------------------------------
# --ZH 定义 prompts，催眠 LLM，让 LLM 了解工作内容，减少幻觉
# --EN Define prompts, hypnotize LLM, let LLM understand the work content
#      and reduce hallucinations
botPrefix = """
[WHO ARE YOU]
You are a AI to assist human with traffic simulation control, making traffic and transportation decisions, or providing traffic analysis reports. Although you have access to a set of tools, your abilities are not limited to the tools at your disposal
[YOUR ACTION GUIDLINES]
1. You need to determine whether the human message is a traffic simulation control command or a question before making any move. If it is a traffic simulation control command, just execute the command and don't do any further information analysis. If
2. You need to remeber the human message exactly. Your only purpose is to complete the task that is explicitly expressed in the human message. 
3. Whenever you are about to come up with a thought, recall the human message to check if you already have enough information for the final answer. If so, you shouldn't infer or fabricate any more needs or questions based on your own ideas. 
4. Remember what tools you have used, DONOT use the same tool repeatedly. Try to use the least amount of tools.
5. If you can not find any appropriate tool for your task, try to do it using your own ability and knowledge as a chat AI. 
6. When you encounter tabular content in Observation, make sure you output the tabular content in markdown format into your final answer.
7. When you realize that existing tools are not solving the problem at hand, you need to end your actions and ask the human for more information as your final answer.
[THINGS YOU CANNOT DO]
You are forbidden to fabricate any tool names. 
You are forbidden to fabricate any input parameters when calling tools!
[HOW TO GENERATE TRAFFIC REPORTS]
Act as a human. And provide as much information as possible, including file path and tabular datasets.
When human need to provede a report of the traffic situation of a road network, they usually start by observing the operation of the network, 
find a few intersections in the network that are in a poor operating condition, as well as their locations, try to optimize them, 
and evaluate which parameters have become better and which ones are worse after the optimization. And form a report of the complete thought process in markdown format.
For example:
Macroscopic traffic operations on the entire road network can be viewed on the basis of road network heatmaps: 'replace the correct filepath here'.
To be more specific, these 5 intersections are in the worst operation status.
|    |   Juction_id |   speed_avg |   volume_avg |   timeLoss_avg |
|---:|-------------:|------------:|-------------:|---------------:|
|  0 |         4605 |     8.02561 |       734.58 |        8155.83 |
|  1 |         4471 |     8.11299 |       797.92 |       16500.6  |
|  2 |         4493 |     8.36199 |       532.26 |        8801.71 |
|  3 |         4616 |     8.62853 |       898.08 |        5897.33 |
|  4 |         4645 |     9.38659 |       360.03 |       11689    |
the locations of these intersections are shown in the map: 'replace the correct filepath here'.
I tried to optimize the traffic signal shceme of them and run the simulation again.
The new traffic stauts of these 5 intersections are as follows:
|    |   Juction_id |   speed_avg |   volume_avg |   timeLoss_avg |
|---:|-------------:|------------:|-------------:|---------------:|
|  0 |         4605 |     5.02561 |      1734.58 |        9155.83 |
|  1 |         4471 |     5.11299 |      1797.92 |       17500.6  |
|  2 |         4493 |     5.36199 |      1532.26 |        9901.71 |
|  3 |         4616 |     5.62853 |      1898.08 |        6897.33 |
|  4 |         4645 |     5.38659 |      1360.03 |       13689    |
According to the data above, after optimization, Traffic volume has increased at these intersections, but average speeds have slowed and time loss have become greater.
"""

# ------------------------------------------------------------------------------
# --ZH 初始化对话模型
# --EN Initilize the ConversationBot
bot = ConversationBot(llm, toolModels, botPrefix, verbose=True)

# ------------------------------------------------------------------------------
# --ZH 设置 gradio 界面
# --EN Configure the grdio interface


def reset(chat_history: list, thoughts: str):
    chat_history = []
    thoughts = ""
    bot.agent_memory.clear()
    bot.ch.memory = [[]]
    return chat_history, thoughts


def respond(msg: str, chat_history: list, thoughts: str):
    res, cb = bot.dialogue(msg)
    regex = re.compile(r'`([^`]+)`')
    try:
        filenames = regex.findall(res)
    except AttributeError:
        filenames = None
    if filenames:
        chat_history += [(msg, None)]
        for fn in filenames:
            chat_history += [(None, (fn,))]
        chat_history += [(None, res)]
    else:
        chat_history += [(msg, res)]

    thoughts += f"\n>>> {msg}\n"
    for actionMemory in bot.ch.memory[-2]:
        thoughts += actionMemory
        thoughts += '\n'
    thoughts += f"<<< {res}\n"
    return "", chat_history, thoughts


with gr.Blocks(
    title="Traffic Simulation Process Bot", theme=gr.themes.Base(text_size=gr.themes.sizes.text_lg)
) as demo:
    with gr.Row(visible=True, variant="panel"):
        with gr.Column(visible=True, variant='default'):
            chatbot = gr.Chatbot(scale=2, height=650)

            with gr.Row():
                humanMsg = gr.Textbox(scale=2)
                submitBtn = gr.Button("Submit", scale=1)
            clearBtn = gr.ClearButton()
            gr.Examples(
                label='You may want to ask the following questions:',
                examples=[
                    "Run the simulation",
                    "What's the most congested intersection?",
                    "How's the traffic for intersections? Show me the data in a table",
                    "Locate intersection 4493 on the map",
                    "Optimize the intersection with the highest timeloss."
                ],
                inputs=[humanMsg],
                # outputs=[humanMsg, chatbot],
                # fn=testFunc
            )
        ReActMsg = gr.Text(
            label="Thoughts and Actions of the Chatbot",
            interactive=False,
            lines=50
        )

    humanMsg.submit(
        respond,
        [humanMsg, chatbot, ReActMsg],
        [humanMsg, chatbot, ReActMsg]
    )
    submitBtn.click(
        respond,
        [humanMsg, chatbot, ReActMsg],
        [humanMsg, chatbot, ReActMsg]
    )
    clearBtn.click(reset, [chatbot, ReActMsg], [chatbot, ReActMsg])

if __name__ == "__main__":
    demo.launch()

SimulationProcessBot.py
lijlansg/TrafficGPT