Skip to content

Commit

Permalink
Allow each member to select its provider, model and temperature
Browse files Browse the repository at this point in the history
  • Loading branch information
@StreetLamb
StreetLamb committed Apr 28, 2024
1 parent 315eb45 commit ac90532
Show file tree
Hide file tree
Showing 13 changed files with 977 additions and 76 deletions.
198 changes: 132 additions & 66 deletions backend/app/core/graph/build.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,38 +4,42 @@

from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph

from app.core.graph.members import (
Leader,
LeaderNode,
Member,
SummariserNode,
Team,
TeamState,
WorkerNode,
)
from app.models import ChatMessage, Team
from app.models import ChatMessage
from app.models import Member as MemberModel
from app.models import Team as TeamModel

model = ChatOpenAI(model="gpt-3.5-turbo")


def convert_team_to_dict(team: Team, members: list[MemberModel]):
"""Convert team and members model to teams dict
def convert_team_to_dict(
team: TeamModel, members: list[MemberModel]
) -> dict[str, Team]:
"""
Converts a team and its members into a dictionary representation.
Args:
team (Team): Team Model
members (list[Member]): list of members model
team (Team): The team model to be converted.
members (list[Member]): A list of member models belonging to the team.
Returns:
dict: A dictionary containing the team's information, with member details.
Raises:
ValueError: Root leader not found
ValueError: If the root leader is not found in the team.
Returns:
dict: Dict containing the team and members
Notes:
This function assumes that each team has a single root leader.
"""

teams = {}
teams: dict[str, Team] = {}

in_counts = defaultdict(int)
out_counts = defaultdict(list[int])
Expand All @@ -61,22 +65,28 @@ def convert_team_to_dict(team: Team, members: list[MemberModel]):
if member.type == "root" or member.type == "leader":
leader_name = member.name
# Create the team definitions
teams[leader_name] = {
"name": team.name,
"members": {},
}
teams[leader_name] = Team(
name=team.name,
model=member.model,
members={},
provider=member.provider,
temperature=member.temperature,
)
# If member is not root team leader, add as a member
if member.type != "root":
member_name = member.name
leader = members_lookup[member.source]
leader_name = leader.name
teams[leader_name]["members"][member_name] = {
"type": member.type,
"name": member_name,
"backstory": member.backstory or "",
"role": member.role,
"tools": [skill.name for skill in member.skills],
}
teams[leader_name].members[member_name] = Member(
type=member.type,
name=member_name,
backstory=member.backstory or "",
role=member.role,
tools=[skill.name for skill in member.skills],
provider=member.provider,
model=member.model,
temperature=member.temperature,
)

for nei_id in out_counts[member_id]:
in_counts[nei_id] -= 1
Expand All @@ -86,16 +96,30 @@ def convert_team_to_dict(team: Team, members: list[MemberModel]):
return teams


# Create the Member/Leader class instance in members
def format_teams(teams: dict[str, any]):
"""Update the team members to use Member/Leader"""
for team in teams:
members = teams[team]["members"]
def format_teams(teams: dict[str, any]) -> dict[str, Team]:
"""
FOR TESTING PURPOSES ONLY!
This function takes a dictionary of teams and formats their member lists to use instances of the `Member` or `Leader`
classes.
Args:
teams (dict[str, any]): A dictionary where each key is a team name and the value is another dictionary containing
the team's members
Returns:
dict[str, Team]: The input dictionary with its member lists formatted to use instances of `Member` or `Leader`
"""
for team_name, team in teams.items():
if not isinstance(team, dict):
raise ValueError(f"Invalid team {team_name}. Teams must be dictionaries.")
members = team.get("members", {})
for k, v in members.items():
teams[team]["members"][k] = (
Leader(**v) if v["type"] == "leader" else Member(**v)
)
return teams
if v["type"] == "leader":
teams[team_name]["members"][k] = Leader(**v)
else:
teams[team_name]["members"][k] = Member(**v)
return {team_name: Team(**team) for team_name, team in teams.items()}


def router(state: TeamState):
Expand Down Expand Up @@ -127,21 +151,56 @@ def exit_chain(state: TeamState):
return {"messages": [answer]}


def create_graph(
teams: dict[str, dict[str, str | dict[str, Member | Leader]]], leader_name: str
):
"""
Create the team's graph.
def create_graph(teams: dict[str, Team], leader_name: str):
"""Create the team's graph.
This function creates a graph representation of the given teams. The graph is represented as a dictionary where each key is a team name,
and the value is another dictionary containing the team's members, their roles, and tools.
Args:
teams (dict[str, dict[str, str | dict[str, Member | Leader]]]): A dictionary where each key is a team leader's name and the value is
another dictionary containing the team's members.
leader_name (str): The name of the root leader in the team.
Returns:
dict: A dictionary representing the graph of teams.
"""
build = StateGraph(TeamState)
# Add the start and end node
build.add_node(leader_name, RunnableLambda(LeaderNode(model).delegate))
build.add_node("summariser", RunnableLambda(SummariserNode(model).summarise))

members = teams[leader_name]["members"]
build.add_node(
leader_name,
RunnableLambda(
LeaderNode(
teams[leader_name].provider,
teams[leader_name].model,
teams[leader_name].temperature,
).delegate
),
)
build.add_node(
"summariser",
RunnableLambda(
SummariserNode(
teams[leader_name].provider,
teams[leader_name].model,
teams[leader_name].temperature,
).summarise
),
)

members = teams[leader_name].members
for name, member in members.items():
if isinstance(member, Member):
build.add_node(name, RunnableLambda(WorkerNode(model).work))
build.add_node(
name,
RunnableLambda(
WorkerNode(
member.provider,
member.model,
member.temperature,
).work
),
)
elif isinstance(member, Leader):
subgraph = create_graph(teams, leader_name=name)
enter = partial(enter_chain, team=teams[name])
Expand All @@ -160,11 +219,12 @@ def create_graph(
return graph


async def generator(team: Team, members: list[Member], messages: list[ChatMessage]):
async def generator(
team: TeamModel, members: list[Member], messages: list[ChatMessage]
):
"""Create the graph and stream responses as JSON."""
teams = convert_team_to_dict(team, members)
team_leader = list(teams.keys())[0]
format_teams(teams)
root = create_graph(teams, leader_name=team_leader)
messages = [
HumanMessage(message.content)
Expand All @@ -177,8 +237,8 @@ async def generator(team: Team, members: list[Member], messages: list[ChatMessag
async for output in root.astream(
{
"messages": messages,
"team_name": teams[team_leader]["name"],
"team_members": teams[team_leader]["members"],
"team_name": teams[team_leader].name,
"team_members": teams[team_leader].members,
}
):
if "__end__" not in output:
Expand All @@ -196,65 +256,71 @@ async def generator(team: Team, members: list[Member], messages: list[ChatMessag
# teams = {
# "FoodExpertLeader": {
# "name": "FoodExperts",
# "model": "ChatOpenAI",
# "members": {
# "ChineseFoodExpert": {
# "type": "worker",
# "name": "ChineseFoodExpert",
# "backstory": "Studied culinary school in Singapore. Well-verse in hawker to fine-dining experiences. ISFP.",
# "role": "Provide chinese food suggestions in Singapore",
# "tools": []
# "tools": [],
# "model": "ChatOpenAI",
# },
# "MalayFoodExpert": {
# "type": "worker",
# "name": "MalayFoodExpert",
# "backstory": "Studied culinary school in Singapore. Well-verse in hawker to fine-dining experiences. INTP.",
# "role": "Provide malay food suggestions in Singapore",
# "tools": []
# "tools": [],
# "model": "ChatOpenAI",
# },
# }
# },
# },
# "TravelExpertLeader": {
# "name": "TravelKakis",
# "model": "ChatOpenAI",
# "members": {
# "FoodExpertLeader": {
# "type": "leader",
# "name": "FoodExpertLeader",
# "role": "Gather inputs from your team and provide a diverse food suggestions in Singapore.",
# "tools": []
# "tools": [],
# "model": "ChatOpenAI",
# },
# "HistoryExpert": {
# "type": "worker",
# "name": "HistoryExpert",
# "backstory": "Studied Singapore history. Well-verse in Singapore architecture. INTJ.",
# "role": "Provide places to sight-see with a history/architecture angle",
# "tools": ["search"]
# }
# }
# }
# "tools": ["search"],
# "model": "ChatOpenAI",
# },
# },
# },
# }

# format_teams(teams)

# teams = format_teams(teams)
# team_leader = "TravelExpertLeader"

# root = create_graph(teams, team_leader)

# messages = [
# HumanMessage(f"What is the best food in Singapore")
# ]
# messages = [HumanMessage("What is the best food in Singapore")]

# initial_state = {
# "messages": messages,
# "team_name": teams[team_leader]["name"],
# "team_members": teams[team_leader]["members"],
# "team_name": teams[team_leader].name,
# "team_members": teams[team_leader].members,
# }


# async def main():
# async for s in root.astream(initial_state):
# if "__end__" not in s:
# print(s)
# print("----")
# print(s)
# print("----")


# import asyncio
# main()
# # import asyncio

# asyncio.run(main())
# # asyncio.run(main())
Loading

0 comments on commit ac90532

Please sign in to comment.