File size: 7,838 Bytes
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
193
194
195
196
197
198
"""
Enrich each team member with a fictional background story and typical job activities.

PROMPT> python -m src.team.enrich_team_members_with_background_story
"""
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."
    )
    job_background_story_of_employee: str = Field(
        description="Provide a fictional story for this person."
    )
    typical_job_activities: str = Field(
        description="Describe some typical activities in the job."
    )

class TeamDetails(BaseModel):
    team_members: list[TeamMember] = Field(
        description="The experts with domain knowledge about the problem."
    )

ENRICH_TEAM_MEMBERS_SYSTEM_PROMPT = """
For each team member provided, enrich them with a fictional background story and typical job activities.

Write a fictional background story about the person. It must be one paragraph that covers: 
- First name and last name.
- Location.
- What education, experience, and skills does this person have.
- Familiarity with the task.
- Why is this particular person is relevant.

The typical_job_activities describes relevant skills needed for this project.
"""

@dataclass
class EnrichTeamMembersWithBackgroundStory:
    """
    Enrich each team member with a fictional background story and typical job activities.
    """
    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]) -> 'EnrichTeamMembersWithBackgroundStory':
        """
        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_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(TeamDetails)
        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_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 = EnrichTeamMembersWithBackgroundStory(
            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_and_merge_with_team_members(team_details: TeamDetails, team_member_list: list[dict]) -> list[dict]:
        result_team_member_list = team_member_list.copy()
        enriched_team_member_list = team_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['typical_job_activities'] = enriched_team_member.typical_job_activities
                team_member['background_story'] = enriched_team_member.job_background_story_of_employee
        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 = "Investigate outbreak of a deadly new disease in the jungle."

    team_member_list = [
        {
            "id": 1,
            "category": "Medical Expertise",
            "explanation": "To identify and understand the pathogen, its symptoms, transmission methods, and potential treatments or vaccines."
        },
        {
            "id": 2,
            "category": "Epidemiologist",
            "explanation": "To track the spread of the disease, analyze patterns, and develop strategies to contain it."
        },
        {
            "id": 3,
            "category": "Field Research Specialist",
            "explanation": "To conduct on-the-ground investigations in challenging jungle environments, collect samples, and interact with affected communities."
        },
        {
            "id": 4,
            "category": "Logistics Coordinator",
            "explanation": "To manage the supply chain for equipment, medical supplies, and personnel transport to remote locations."
        }
    ]

    query = EnrichTeamMembersWithBackgroundStory.format_query(job_description, team_member_list)
    print(f"Query:\n{query}\n\n")

    enrich_team_members_with_background_story = EnrichTeamMembersWithBackgroundStory.execute(llm, query, team_member_list)
    json_response = enrich_team_members_with_background_story.to_dict(include_system_prompt=False, include_user_prompt=False)
    print(json.dumps(json_response, indent=2))

    print("\n\nTeam members with extra details:")
    enriched_json = enrich_team_members_with_background_story.team_member_list
    print(json.dumps(enriched_json, indent=2))