File size: 7,602 Bytes
6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 6369972 f0808c8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
"""
Enrich the team members with what kind of equipment and facilities they need for the task.
PROMPT> python -m src.team.enrich_team_members_with_environment_info
"""
import os
import json
import time
import logging
from math import ceil
from dataclasses import dataclass
from typing import List, Optional
from pydantic import BaseModel, Field
from llama_index.core.llms import ChatMessage, MessageRole
from llama_index.core.llms.llm import LLM
from src.format_json_for_use_in_query import format_json_for_use_in_query
logger = logging.getLogger(__name__)
class TeamMember(BaseModel):
"""A human with domain knowledge."""
id: int = Field(
description="A unique id for the job_category."
)
equipment_needs: str = Field(
description="What expensive resources are needed for the daily job."
)
facility_needs: str = Field(
description="What facilities are needed for the daily job."
)
class DocumentDetails(BaseModel):
team_members: list[TeamMember] = Field(
description="The experts with domain knowledge about the problem."
)
ENRICH_TEAM_MEMBERS_ENVIRONMENT_INFO_SYSTEM_PROMPT = """
You are an expert at determining what equipment and facilities are needed for different job roles given a project description.
For each team member provided, identify the specific equipment and facilities they require to effectively perform their daily tasks within the context of the given project description. Provide concise but descriptive answers.
"""
@dataclass
class EnrichTeamMembersWithEnvironmentInfo:
"""
Enrich each team member with more info.
"""
system_prompt: str
user_prompt: str
response: dict
metadata: dict
team_member_list: list[dict]
@classmethod
def format_query(cls, job_description: str, team_member_list: list[dict]) -> str:
if not isinstance(job_description, str):
raise ValueError("Invalid job_description.")
if not isinstance(team_member_list, list):
raise ValueError("Invalid team_member_list.")
query = (
f"Project description:\n{job_description}\n\n"
f"Here is the list of team members that needs to be enriched:\n{format_json_for_use_in_query(team_member_list)}"
)
return query
@classmethod
def execute(cls, llm: LLM, user_prompt: str, team_member_list: list[dict]) -> 'EnrichTeamMembersWithEnvironmentInfo':
"""
Invoke LLM with each team member.
"""
if not isinstance(llm, LLM):
raise ValueError("Invalid LLM instance.")
if not isinstance(user_prompt, str):
raise ValueError("Invalid user_prompt.")
if not isinstance(team_member_list, list):
raise ValueError("Invalid team_member_list.")
logger.debug(f"User Prompt:\n{user_prompt}")
system_prompt = ENRICH_TEAM_MEMBERS_ENVIRONMENT_INFO_SYSTEM_PROMPT.strip()
chat_message_list = [
ChatMessage(
role=MessageRole.SYSTEM,
content=system_prompt,
),
ChatMessage(
role=MessageRole.USER,
content=user_prompt,
)
]
sllm = llm.as_structured_llm(DocumentDetails)
start_time = time.perf_counter()
try:
chat_response = sllm.chat(chat_message_list)
except Exception as e:
logger.debug(f"LLM chat interaction failed: {e}")
logger.error("LLM chat interaction failed.", exc_info=True)
raise ValueError("LLM chat interaction failed.") from e
end_time = time.perf_counter()
duration = int(ceil(end_time - start_time))
response_byte_count = len(chat_response.message.content.encode('utf-8'))
logger.info(f"LLM chat interaction completed in {duration} seconds. Response byte count: {response_byte_count}")
json_response = chat_response.raw.model_dump()
team_member_list_enriched = cls.cleanup_enriched_team_members_with_environment_info_and_merge_with_team_members(chat_response.raw, team_member_list)
metadata = dict(llm.metadata)
metadata["llm_classname"] = llm.class_name()
metadata["duration"] = duration
metadata["response_byte_count"] = response_byte_count
result = EnrichTeamMembersWithEnvironmentInfo(
system_prompt=system_prompt,
user_prompt=user_prompt,
response=json_response,
metadata=metadata,
team_member_list=team_member_list_enriched,
)
return result
def to_dict(self, include_metadata=True, include_system_prompt=True, include_user_prompt=True) -> dict:
d = self.response.copy()
if include_metadata:
d['metadata'] = self.metadata
if include_system_prompt:
d['system_prompt'] = self.system_prompt
if include_user_prompt:
d['user_prompt'] = self.user_prompt
return d
def cleanup_enriched_team_members_with_environment_info_and_merge_with_team_members(document_details: DocumentDetails, team_member_list: list[dict]) -> list:
result_team_member_list = team_member_list.copy()
enriched_team_member_list = document_details.team_members
id_to_enriched_team_member = {item.id: item for item in enriched_team_member_list}
for team_member_index, team_member in enumerate(result_team_member_list):
if not 'id' in team_member:
logger.warning(f"Team member #{team_member_index} does not have an id")
continue
id = team_member['id']
enriched_team_member = id_to_enriched_team_member.get(id)
if enriched_team_member:
team_member['equipment_needs'] = enriched_team_member.equipment_needs
team_member['facility_needs'] = enriched_team_member.facility_needs
return result_team_member_list
if __name__ == "__main__":
from src.llm_factory import get_llm
llm = get_llm("ollama-llama3.1")
# llm = get_llm("deepseek-chat")
job_description = "Establish a new police station in a high crime area."
team_member_list = [
{
"id": 1,
"category": "Law Enforcement",
"explanation": "Police officers and detectives are essential for patrolling, investigation, and maintaining public safety."
},
{
"id": 2,
"category": "Administration",
"explanation": "Administrative staff manage paperwork, scheduling, and coordination of police activities."
},
{
"id": 3,
"category": "Forensics",
"explanation": "Forensic experts analyze crime scene evidence to support investigations."
},
{
"id": 4,
"category": "Community Relations",
"explanation": "Officers or liaisons engage with the community to build trust and cooperation."
}
]
query = EnrichTeamMembersWithEnvironmentInfo.format_query(job_description, team_member_list)
print(f"Query:\n{query}\n\n")
enrich_team_members_with_environment_info = EnrichTeamMembersWithEnvironmentInfo.execute(llm, query, team_member_list)
json_response = enrich_team_members_with_environment_info.to_dict(include_system_prompt=False, include_user_prompt=False)
print(json.dumps(json_response, indent=2))
print("\n\nTeam members:")
json_team = enrich_team_members_with_environment_info.team_member_list
print(json.dumps(json_team, indent=2))
|