Skip to content

Commit f842887

Browse files
StreetLambStreetLamb
authored
Merge pull request #1 from StreetLamb/backend/langchain
Setup endpoints and graph module to build and run graph
2 parents 7b6b092 + 65c1fef commit f842887

File tree

9 files changed

+1771
-21
lines changed

9 files changed

+1771
-21
lines changed

‎backend/app/api/main.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
from fastapi import APIRouter
22

3-
from app.api.routes import items, login, users, utils
3+
from app.api.routes import items, login, users, utils, teams
44

55
api_router = APIRouter()
66
api_router.include_router(login.router, tags=["login"])
77
api_router.include_router(users.router, prefix="/users", tags=["users"])
88
api_router.include_router(utils.router, prefix="/utils", tags=["utils"])
99
api_router.include_router(items.router, prefix="/items", tags=["items"])
10+
api_router.include_router(teams.router, prefix="/teams", tags=["teams"])

‎backend/app/api/routes/teams.py

+152
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
from typing import Any
2+
from sqlmodel import func, select
3+
from fastapi import APIRouter, HTTPException
4+
from fastapi.responses import StreamingResponse
5+
6+
from app.core.graph.build import generator
7+
from app.api.deps import CurrentUser, SessionDep
8+
from app.models import TeamChat, TeamsOut, TeamCreate, TeamOut, Team, Message
9+
10+
# TODO: To remove
11+
teams = {
12+
"FoodExpertLeader": {
13+
"name": "FoodExperts",
14+
"members": {
15+
"ChineseFoodExpert": {
16+
"type": "worker",
17+
"name": "ChineseFoodExpert",
18+
"backstory": "Studied culinary school in Singapore. Well-verse in hawker to fine-dining experiences. ISFP.",
19+
"role": "Provide chinese food suggestions in Singapore",
20+
"tools": []
21+
},
22+
"MalayFoodExpert": {
23+
"type": "worker",
24+
"name": "MalayFoodExpert",
25+
"backstory": "Studied culinary school in Singapore. Well-verse in hawker to fine-dining experiences. INTP.",
26+
"role": "Provide malay food suggestions in Singapore",
27+
"tools": []
28+
},
29+
}
30+
},
31+
"TravelExpertLeader": {
32+
"name": "TravelKakis",
33+
"members": {
34+
"FoodExpertLeader": {
35+
"type": "leader",
36+
"name": "FoodExpertLeader",
37+
"role": "Gather inputs from your team and provide a diverse food suggestions in Singapore.",
38+
"tools": []
39+
},
40+
"HistoryExpert": {
41+
"type": "worker",
42+
"name": "HistoryExpert",
43+
"backstory": "Studied Singapore history. Well-verse in Singapore architecture. INTJ.",
44+
"role": "Provide places to sight-see with a history/architecture angle",
45+
"tools": []
46+
}
47+
}
48+
}
49+
}
50+
team_leader = "TravelExpertLeader"
51+
52+
router = APIRouter()
53+
54+
@router.get("/", response_model=TeamsOut)
55+
def read_teams(
56+
session: SessionDep, current_user: CurrentUser, skip: int = 0, limit: int = 100
57+
) -> Any:
58+
"""
59+
Retrieve teams
60+
"""
61+
62+
if current_user.is_superuser:
63+
count_statement = select(func.count()).select_from(Team)
64+
count = session.exec(count_statement).one()
65+
statement = select(Team).offset(skip).limit(limit)
66+
teams = session.exec(statement).all()
67+
else:
68+
count_statement = (
69+
select(func.count())
70+
.select_from(Team)
71+
.where(Team.owner_id == current_user.id)
72+
)
73+
count = session.exec(count_statement).one()
74+
statement = (
75+
select(Team)
76+
.where(Team.owner_id == current_user.id)
77+
.offset(skip)
78+
.limit(limit)
79+
)
80+
teams = session.exec(statement).all()
81+
return TeamsOut(data=teams, count=count)
82+
83+
@router.get("/{id}", response_model=TeamOut)
84+
def read_team(session: SessionDep, current_user: CurrentUser, id: int) -> Any:
85+
"""
86+
Get team by ID.
87+
"""
88+
team = session.get(Team, id)
89+
if not team:
90+
raise HTTPException(status_code=404, detail="Team not found")
91+
if not current_user.is_superuser and (team.owner_id != current_user.id):
92+
raise HTTPException(status_code=400, detail="Not enough permissions")
93+
return team
94+
95+
@router.post("/", response_model=TeamOut)
96+
def create_team(
97+
*, session: SessionDep, current_user: CurrentUser, team_in: TeamCreate
98+
) -> Any:
99+
"""
100+
Create new team.
101+
"""
102+
team = Team.model_validate(team_in, update={"owner_id": current_user.id})
103+
session.add(team)
104+
session.commit()
105+
session.refresh(team)
106+
return team
107+
108+
@router.put("/{id}", response_model=TeamOut)
109+
def update_team(
110+
*, session: SessionDep, current_user: CurrentUser, id: int, team_in: TeamCreate
111+
) -> Any:
112+
"""
113+
Update a team.
114+
"""
115+
team = session.get(Team, id)
116+
if not team:
117+
raise HTTPException(status_code=404, detail="Team not found")
118+
if not current_user.is_superuser and (team.owner_id != current_user.id):
119+
raise HTTPException(status_code=400, detail="Not enough permissions")
120+
update_dict = team_in.model_dump(exclude_unset=True)
121+
team.sqlmodel_update(update_dict)
122+
session.add(team)
123+
session.commit()
124+
session.refresh(team)
125+
return team
126+
127+
@router.delete("/{id}")
128+
def delete_team(session: SessionDep, current_user: CurrentUser, id: int) -> Any:
129+
"""
130+
Delete a team.
131+
"""
132+
team = session.get(Team, id)
133+
if not team:
134+
raise HTTPException(status_code=404, detail="Team not found")
135+
if not current_user.is_superuser and (team.owner_id != current_user.id):
136+
raise HTTPException(status_code=400, detail="Not enough permissions")
137+
session.delete(team)
138+
session.commit()
139+
return Message(message="Team deleted successfully")
140+
141+
@router.post("/{id}/stream")
142+
async def stream(session: SessionDep, current_user: CurrentUser, id: int, team_chat: TeamChat):
143+
"""
144+
Stream a response to a user's input.
145+
"""
146+
team = session.get(Team, id)
147+
if not team:
148+
raise HTTPException(status_code=404, detail="Team not found")
149+
if not current_user.is_superuser and (team.owner_id != current_user.id):
150+
raise HTTPException(status_code=400, detail="Not enough permissions")
151+
152+
return StreamingResponse(generator(teams, team_leader, team_chat.messages), media_type="text/event-stream")

‎backend/app/core/graph/__init__.py

Whitespace-only changes.

‎backend/app/core/graph/build.py

+162
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
from functools import partial
2+
from typing import Dict, List
3+
from app.models import ChatMessage
4+
from langgraph.graph import StateGraph, END
5+
from langchain_openai import ChatOpenAI
6+
from app.core.graph.members import Leader, LeaderNode, Member, SummariserNode, TeamState, WorkerNode
7+
from langchain_core.messages import HumanMessage, AIMessage
8+
from langchain_core.runnables import RunnableLambda
9+
10+
model = ChatOpenAI(model="gpt-3.5-turbo")
11+
12+
# Create the Member/Leader class instance in members
13+
def format_teams(teams: Dict[str, any]):
14+
"""Update the team members to use Member/Leader"""
15+
for team in teams:
16+
members = teams[team]["members"]
17+
for k,v in members.items():
18+
print(v)
19+
teams[team]["members"][k] = Leader(**v) if v["type"] == "leader" else Member(**v)
20+
return teams
21+
22+
def router(state: TeamState):
23+
return state["next"]
24+
25+
def enter_chain(state: TeamState, team: Dict[str, str | List[Member | Leader]]):
26+
"""
27+
Initialise the sub-graph state.
28+
This makes it so that the states of each graph don't get intermixed.
29+
"""
30+
task = state["task"]
31+
team_name = team["name"]
32+
team_members = team["members"]
33+
34+
results = {
35+
"messages": task,
36+
"team_name": team_name,
37+
"team_members": team_members,
38+
}
39+
return results
40+
41+
def exit_chain(state: TeamState):
42+
"""
43+
Pass the final response back to the top-level graph's state.
44+
"""
45+
answer = state["messages"][-1]
46+
return {"messages": [answer]}
47+
48+
def create_graph(teams: Dict[str, Dict[str, str | Dict[str, Member | Leader]]], leader_name: str):
49+
"""
50+
Create the team's graph.
51+
"""
52+
build = StateGraph(TeamState)
53+
# Add the start and end node
54+
build.add_node(leader_name, RunnableLambda(LeaderNode(model).delegate))
55+
build.add_node("summariser", RunnableLambda(SummariserNode(model).summarise))
56+
57+
members = teams[leader_name]["members"]
58+
for name, member in members.items():
59+
if isinstance(member, Member):
60+
build.add_node(name, RunnableLambda(WorkerNode(model).work))
61+
elif isinstance(member, Leader):
62+
subgraph = create_graph(teams, leader_name=name)
63+
enter = partial(enter_chain, team=teams[name])
64+
build.add_node(name, enter | subgraph | exit_chain)
65+
else:
66+
continue
67+
build.add_edge(name, leader_name)
68+
69+
conditional_mapping = {v:v for v in members}
70+
conditional_mapping["FINISH"] = "summariser"
71+
build.add_conditional_edges(leader_name, router, conditional_mapping)
72+
73+
build.set_entry_point(leader_name)
74+
build.set_finish_point("summariser")
75+
graph = build.compile()
76+
return graph
77+
78+
79+
80+
async def generator(teams: dict, team_leader: str, messages: List[ChatMessage]):
81+
"""Create the graph and strem the response"""
82+
format_teams(teams)
83+
root = create_graph(teams, leader_name=team_leader)
84+
messages = [HumanMessage(message.content) if message.type == "human" else AIMessage(message.content) for message in messages]
85+
86+
async for output in root.astream({
87+
"messages": messages,
88+
"team_name": teams[team_leader]["name"],
89+
"team_members": teams[team_leader]["members"]
90+
}):
91+
for key, value in output.items():
92+
if key != "__end__":
93+
response = {key :value}
94+
formatted_output = f"data: {response}\n\n"
95+
print(formatted_output)
96+
yield formatted_output
97+
98+
# teams = {
99+
# "FoodExpertLeader": {
100+
# "name": "FoodExperts",
101+
# "members": {
102+
# "ChineseFoodExpert": {
103+
# "type": "worker",
104+
# "name": "ChineseFoodExpert",
105+
# "backstory": "Studied culinary school in Singapore. Well-verse in hawker to fine-dining experiences. ISFP.",
106+
# "role": "Provide chinese food suggestions in Singapore",
107+
# "tools": []
108+
# },
109+
# "MalayFoodExpert": {
110+
# "type": "worker",
111+
# "name": "MalayFoodExpert",
112+
# "backstory": "Studied culinary school in Singapore. Well-verse in hawker to fine-dining experiences. INTP.",
113+
# "role": "Provide malay food suggestions in Singapore",
114+
# "tools": []
115+
# },
116+
# }
117+
# },
118+
# "TravelExpertLeader": {
119+
# "name": "TravelKakis",
120+
# "members": {
121+
# "FoodExpertLeader": {
122+
# "type": "leader",
123+
# "name": "FoodExpertLeader",
124+
# "role": "Gather inputs from your team and provide a diverse food suggestions in Singapore.",
125+
# "tools": []
126+
# },
127+
# "HistoryExpert": {
128+
# "type": "worker",
129+
# "name": "HistoryExpert",
130+
# "backstory": "Studied Singapore history. Well-verse in Singapore architecture. INTJ.",
131+
# "role": "Provide places to sight-see with a history/architecture angle",
132+
# "tools": ["search"]
133+
# }
134+
# }
135+
# }
136+
# }
137+
138+
# format_teams(teams)
139+
140+
# team_leader = "TravelExpertLeader"
141+
142+
# root = create_graph(teams, team_leader)
143+
144+
# messages = [
145+
# HumanMessage(f"What is the best food in Singapore")
146+
# ]
147+
148+
# initial_state = {
149+
# "messages": messages,
150+
# "team_name": teams[team_leader]["name"],
151+
# "team_members": teams[team_leader]["members"],
152+
# }
153+
154+
# async def main():
155+
# async for s in root.astream(initial_state):
156+
# if "__end__" not in s:
157+
# print(s)
158+
# print("----")
159+
160+
# import asyncio
161+
162+
# asyncio.run(main())

0 commit comments

Comments
 (0)