File size: 8,844 Bytes
971b0ff
 
 
 
 
cf6b61d
 
 
 
 
 
971b0ff
cf6b61d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
971b0ff
 
 
 
 
 
 
 
 
 
 
cf6b61d
971b0ff
 
 
 
 
cf6b61d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
971b0ff
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cf6b61d
 
971b0ff
cf6b61d
971b0ff
 
cf6b61d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
971b0ff
cf6b61d
 
 
 
 
 
971b0ff
cf6b61d
971b0ff
 
 
 
 
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
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
# chat.py
# Filepath: ai-email-assistant/chat.py
# Implements the chat interface and message handling

import gradio as gr
import os
import json
import logging
from groq import Groq
from dotenv import load_dotenv

from web_search import search_web
from google_sheets_utils import (
    get_or_create_spreadsheet,
    get_profiles_from_sheet,
    delete_profile_from_sheet,
    get_profile_summaries,
    get_template_summaries,
    get_profile_by_id,
    get_template_by_id,
    save_profile,
    update_profile_in_sheet,
    delete_template_from_sheet,
    save_template,
    update_template,
    generate_template_id,
    generate_profile_id
)
from email_actions import get_templates

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def load_system_prompt(filename):
    try:
        with open(filename, 'r') as file:
            return file.read().strip()
    except FileNotFoundError:
        return "System prompt file not found."

def create_chat_interface(sender_summaries, receiver_summaries, template_summaries,
                          current_sender, current_receiver, current_template,
                          sender_components, receiver_components, email_actions_components):
    system_prompt = load_system_prompt('system_prompt.txt')
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    # Global spreadsheet variable
    global spreadsheet
    spreadsheet = get_or_create_spreadsheet()
    
    # Initialize Groq client
    load_dotenv()  # Load environment variables from .env file
    groq_api_key = os.getenv('GROQ_API_KEY')
    if not groq_api_key:
        logger.error("GROQ_API_KEY not found in environment variables")
        raise ValueError("GROQ_API_KEY not found in environment variables")
    client = Groq(api_key=groq_api_key)
    model = "llama3-groq-70b-8192-tool-use-preview"

    # Define available functions (tools) for the assistant
    def get_profile_summaries_tool(profile_type: str):
        """Get summaries of profiles (Sender or Receiver)"""
        return get_profile_summaries(spreadsheet.worksheet(profile_type), profile_type)

    def get_template_summaries_tool():
        """Get summaries of email templates"""
        return get_template_summaries(spreadsheet)

    def get_profile_by_id_tool(profile_type: str, profile_id: str):
        """Get a specific profile by ID"""
        return get_profile_by_id(spreadsheet.worksheet(profile_type), profile_id)

    def get_template_by_id_tool(template_id: str):
        """Get a specific template by ID"""
        return get_template_by_id(spreadsheet.worksheet("Email Templates"), template_id)

    def create_profile_tool(profile_type: str, name: str, email: str, position: str, company: str, context: str):
        """Create a new profile"""
        return save_profile(spreadsheet.worksheet(profile_type), [profile_type, generate_profile_id(), name, email, position, company, context])

    def update_profile_tool(profile_type: str, profile_id: str, name: str, email: str, position: str, company: str, context: str):
        """Update an existing profile"""
        return update_profile_in_sheet(spreadsheet.worksheet(profile_type), profile_id, [profile_type, profile_id, name, email, position, company, context])

    def delete_profile_tool(profile_type: str, profile_id: str):
        """Delete a profile"""
        return delete_profile_from_sheet(spreadsheet.worksheet(profile_type), profile_id)

    def create_template_tool(name: str, template_type: str, subject: str, body: str):
        """Create a new email template"""
        return save_template(spreadsheet.worksheet("Email Templates"), generate_template_id(), name, template_type, subject, body)

    def update_template_tool(template_id: str, name: str, template_type: str, subject: str, body: str):
        """Update an existing email template"""
        return update_template(spreadsheet.worksheet("Email Templates"), template_id, name, template_type, subject, body)

    def delete_template_tool(template_id: str):
        """Delete an email template"""
        return delete_template_from_sheet(spreadsheet.worksheet("Email Templates"), template_id)

    # Map function names to actual functions
    available_functions = {
        "get_profile_summaries": get_profile_summaries_tool,
        "get_template_summaries": get_template_summaries_tool,
        "get_profile_by_id": get_profile_by_id_tool,
        "get_template_by_id": get_template_by_id_tool,
        "create_profile": create_profile_tool,
        "update_profile": update_profile_tool,
        "delete_profile": delete_profile_tool,
        "create_template": create_template_tool,
        "update_template": update_template_tool,
        "delete_template": delete_template_tool,
        "search_web": search_web,
    }

    # Define functions (tools) for the assistant to use
    tools = []
    for func_name, func in available_functions.items():
        tool = {
            "type": "function",
            "function": {
                "name": func_name,
                "description": func.__doc__ or "",
                "parameters": {
                    "type": "object",
                    "properties": {},  # Parameters will be auto-inferred or can be explicitly defined
                    "required": []
                },
            },
        }
        tools.append(tool)

    def get_context():
        context = "Current context:\n"
        
        if current_sender.value:
            context += f"Current Sender: {current_sender.value}\n"
        if current_receiver.value:
            context += f"Current Receiver: {current_receiver.value}\n"
        if current_template.value:
            context += f"Current Template: {current_template.value}\n"
        
        context += f"Available Sender Profiles: {', '.join(sender_summaries)}\n"
        context += f"Available Receiver Profiles: {', '.join(receiver_summaries)}\n"
        context += f"Available Templates: {', '.join(template_summaries)}\n"
        
        return context

    # Initialize chat history
    chat_history = []

    def respond(user_input, history):
        context = get_context()
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "system", "content": context}
        ]
        # Append conversation history
        for human_msg, ai_msg in history:
            messages.append({"role": "user", "content": human_msg})
            messages.append({"role": "assistant", "content": ai_msg})
        # Append latest user input
        messages.append({"role": "user", "content": user_input})

        response = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=tools,
            tool_choice="auto",
            max_tokens=512  # Adjust as needed
        )

        response_message = response.choices[0].message

        # Check if the assistant made any tool calls
        if hasattr(response_message, 'tool_calls') and response_message.tool_calls:
            # Process each tool call
            for tool_call in response_message.tool_calls:
                function_name = tool_call.function.name
                function_args = json.loads(tool_call.function.arguments)
                function_to_call = available_functions.get(function_name)

                if function_to_call:
                    # Call the function with arguments
                    function_response = function_to_call(**function_args)

                    # Add the function response to the messages
                    messages.append({
                        "role": "tool",
                        "content": json.dumps(function_response),
                        "tool_call_id": tool_call.id,
                    })
                else:
                    # If function not found, add an error message
                    messages.append({
                        "role": "tool",
                        "content": json.dumps({"error": f"Function {function_name} not found."}),
                        "tool_call_id": tool_call.id,
                    })

            # Send the updated messages back to the model for final response
            response = client.chat.completions.create(
                model=model,
                messages=messages,
                tools=tools,
                tool_choice="auto",
                max_tokens=512
            )
            response_message = response.choices[0].message

        assistant_response = response_message.content

        # Update chat history
        history.append((user_input, assistant_response))

        return "", history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

    return chatbot, msg, clear