-
Notifications
You must be signed in to change notification settings - Fork 699
/
main.py
415 lines (365 loc) · 16.5 KB
/
main.py
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
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
from anthropic import Anthropic
import os
from datetime import datetime
import json
from colorama import init, Fore, Style
from pygments import highlight
from pygments.lexers import get_lexer_by_name
from pygments.formatters import TerminalFormatter
from tavily import TavilyClient
import pygments.util
import base64
from PIL import Image
import io
# Initialize colorama
init()
# Color constants
USER_COLOR = Fore.WHITE
CLAUDE_COLOR = Fore.BLUE
TOOL_COLOR = Fore.YELLOW
RESULT_COLOR = Fore.GREEN
# Initialize the Anthropic client
client = Anthropic(api_key="YOUR_API_KEY")
# Initialize the Tavily client
tavily = TavilyClient(api_key="YOUR_API_KEY")
# Set up the conversation memory
conversation_history = []
# System prompt
system_prompt = """
You are Claude, an AI assistant powered by Anthropic's Claude-3.5-Sonnet model. You are an exceptional software developer with vast knowledge across multiple programming languages, frameworks, and best practices. Your capabilities include:
1. Creating project structures, including folders and files
2. Writing clean, efficient, and well-documented code
3. Debugging complex issues and providing detailed explanations
4. Offering architectural insights and design patterns
5. Staying up-to-date with the latest technologies and industry trends
6. Reading and analyzing existing files in the project directory
7. Listing files in the root directory of the project
8. Performing web searches to get up-to-date information or additional context
9. When you use search make sure you use the best query to get the most accurate and up-to-date information
10. You NEVER remove existing code if doesnt require to be changed or removed, never use comments like # ... (keep existing code) ... or # ... (rest of the code) ... etc, you only add new code or remove it.
11. Analyzing images provided by the user
When an image is provided, carefully analyze its contents and incorporate your observations into your responses.
When asked to create a project:
- Always start by creating a root folder for the project.
- Then, create the necessary subdirectories and files within that root folder.
- Organize the project structure logically and follow best practices for the specific type of project being created.
- Use the provided tools to create folders and files as needed.
When asked to make edits or improvements:
- Use the read_file tool to examine the contents of existing files.
- Analyze the code and suggest improvements or make necessary edits.
- Use the write_to_file tool to implement changes.
Be sure to consider the type of project (e.g., Python, JavaScript, web application) when determining the appropriate structure and files to include.
You can now read files, list the contents of the root folder where this script is being run, and perform web searches. Use these capabilities when:
- The user asks for edits or improvements to existing files
- You need to understand the current state of the project
- You believe reading a file or listing directory contents will be beneficial to accomplish the user's goal
- You need up-to-date information or additional context to answer a question accurately
When you need current information or feel that a search could provide a better answer, use the tavily_search tool. This tool performs a web search and returns a concise answer along with relevant sources.
Always strive to provide the most accurate, helpful, and detailed responses possible. If you're unsure about something, admit it and consider using the search tool to find the most current information.
Answer the user's request using relevant tools (if they are available). Before calling a tool, do some analysis within \<thinking>\</thinking> tags. First, think about which of the provided tools is the relevant tool to answer the user's request. Second, go through each of the required parameters of the relevant tool and determine if the user has directly provided or given enough information to infer a value. When deciding if the parameter can be inferred, carefully consider all the context to see if it supports a specific value. If all of the required parameters are present or can be reasonably inferred, close the thinking tag and proceed with the tool call. BUT, if one of the values for a required parameter is missing, DO NOT invoke the function (not even with fillers for the missing params) and instead, ask the user to provide the missing parameters. DO NOT ask for more information on optional parameters if it is not provided.
"""
# Helper function to print colored text
def print_colored(text, color):
print(f"{color}{text}{Style.RESET_ALL}")
# Helper function to format and print code
def print_code(code, language):
try:
lexer = get_lexer_by_name(language, stripall=True)
formatted_code = highlight(code, lexer, TerminalFormatter())
print(formatted_code)
except pygments.util.ClassNotFound:
# If the language is not recognized, fall back to plain text
print_colored(f"Code (language: {language}):\n{code}", CLAUDE_COLOR)
# Function to create a folder
def create_folder(path):
try:
os.makedirs(path, exist_ok=True)
return f"Folder created: {path}"
except Exception as e:
return f"Error creating folder: {str(e)}"
# Function to create a file
def create_file(path, content=""):
try:
with open(path, 'w') as f:
f.write(content)
return f"File created: {path}"
except Exception as e:
return f"Error creating file: {str(e)}"
# Function to write to a file
def write_to_file(path, content):
try:
with open(path, 'w') as f:
f.write(content)
return f"Content written to file: {path}"
except Exception as e:
return f"Error writing to file: {str(e)}"
# Function to read a file
def read_file(path):
try:
with open(path, 'r') as f:
content = f.read()
return content
except Exception as e:
return f"Error reading file: {str(e)}"
# Function to list files in the root directory
def list_files(path="."):
try:
files = os.listdir(path)
return "\n".join(files)
except Exception as e:
return f"Error listing files: {str(e)}"
# Function to perform a Tavily search
def tavily_search(query):
try:
response = tavily.qna_search(query=query, search_depth="advanced")
return response
except Exception as e:
return f"Error performing search: {str(e)}"
# Define the tools
tools = [
{
"name": "create_folder",
"description": "Create a new folder at the specified path. Use this when you need to create a new directory in the project structure.",
"input_schema": {
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "The path where the folder should be created"
}
},
"required": ["path"]
}
},
{
"name": "create_file",
"description": "Create a new file at the specified path with optional content. Use this when you need to create a new file in the project structure.",
"input_schema": {
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "The path where the file should be created"
},
"content": {
"type": "string",
"description": "The initial content of the file (optional)"
}
},
"required": ["path"]
}
},
{
"name": "write_to_file",
"description": "Write content to an existing file at the specified path. Use this when you need to add or update content in an existing file.",
"input_schema": {
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "The path of the file to write to"
},
"content": {
"type": "string",
"description": "The content to write to the file"
}
},
"required": ["path", "content"]
}
},
{
"name": "read_file",
"description": "Read the contents of a file at the specified path. Use this when you need to examine the contents of an existing file.",
"input_schema": {
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "The path of the file to read"
}
},
"required": ["path"]
}
},
{
"name": "list_files",
"description": "List all files and directories in the root folder where the script is running. Use this when you need to see the contents of the current directory.",
"input_schema": {
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "The path of the folder to list (default: current directory)"
}
}
}
},
{
"name": "tavily_search",
"description": "Perform a web search using Tavily API to get up-to-date information or additional context. Use this when you need current information or feel a search could provide a better answer.",
"input_schema": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query"
}
},
"required": ["query"]
}
}
]
# Function to execute tools
def execute_tool(tool_name, tool_input):
if tool_name == "create_folder":
return create_folder(tool_input["path"])
elif tool_name == "create_file":
return create_file(tool_input["path"], tool_input.get("content", ""))
elif tool_name == "write_to_file":
return write_to_file(tool_input["path"], tool_input.get("content", ""))
elif tool_name == "read_file":
return read_file(tool_input["path"])
elif tool_name == "list_files":
return list_files(tool_input.get("path", "."))
elif tool_name == "tavily_search":
return tavily_search(tool_input["query"])
else:
return f"Unknown tool: {tool_name}"
# Function to encode image to base64
def encode_image_to_base64(image_path):
try:
with Image.open(image_path) as img:
# Resize image if it's too large
max_size = (1024, 1024)
img.thumbnail(max_size, Image.DEFAULT_STRATEGY)
# Convert image to RGB if it's not
if img.mode != 'RGB':
img = img.convert('RGB')
img_byte_arr = io.BytesIO()
img.save(img_byte_arr, format='JPEG')
return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
except Exception as e:
return f"Error encoding image: {str(e)}"
# Function to send a message to Claude and get the response
def chat_with_claude(user_input, image_path=None):
global conversation_history
if image_path:
# Encode image to base64
image_base64 = encode_image_to_base64(image_path)
#print(f"Image encoded to base64: {image_base64}")
#exit()
# Prepare the message with image
image_message = {
"role": "user",
"content": [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/jpeg",
"data": image_base64
}
},
{
"type": "text",
"text": f"User input for image: {user_input}"
}
]
}
conversation_history.append(image_message)
else:
# Add regular user input to conversation history
conversation_history.append({"role": "user", "content": user_input})
# Prepare the messages for the API call
messages = [msg for msg in conversation_history if msg.get('content')]
# Make the initial API call
response = client.messages.create(
model="claude-3-5-sonnet-20240620",
max_tokens=4000,
system=system_prompt,
messages=messages,
tools=tools,
tool_choice={"type": "auto"}
)
assistant_response = ""
# Process the response
for content_block in response.content:
if content_block.type == "text":
assistant_response += content_block.text
print_colored(f"\nClaude: {content_block.text}", CLAUDE_COLOR)
elif content_block.type == "tool_use":
tool_name = content_block.name
tool_input = content_block.input
tool_use_id = content_block.id
print_colored(f"\nTool Used: {tool_name}", TOOL_COLOR)
print_colored(f"Tool Input: {tool_input}", TOOL_COLOR)
# Execute the tool
result = execute_tool(tool_name, tool_input)
print_colored(f"Tool Result: {result}", RESULT_COLOR)
# Add tool use and result to conversation history
conversation_history.append({"role": "assistant", "content": [content_block]})
conversation_history.append({
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": tool_use_id,
"content": result
}
]
})
# Make another API call with the updated conversation history
tool_response = client.messages.create(
model="claude-3-5-sonnet-20240620",
max_tokens=4000,
system=system_prompt,
messages=[msg for msg in conversation_history if msg.get('content')],
tools=tools,
tool_choice={"type": "auto"}
)
# Process the tool response
for tool_content_block in tool_response.content:
if tool_content_block.type == "text":
assistant_response += tool_content_block.text
print_colored(f"\nClaude: {tool_content_block.text}", CLAUDE_COLOR)
# Add final assistant response to conversation history
if assistant_response:
conversation_history.append({"role": "assistant", "content": assistant_response})
return assistant_response
# Main chat loop
def main():
print_colored("Welcome to the Claude-3.5-Sonnet Engineer Chat with Image Support!", CLAUDE_COLOR)
print_colored("Type 'exit' to end the conversation.", CLAUDE_COLOR)
print_colored("To include an image, drag and drop it into the terminal and press enter. Provide your prompt on next line.", CLAUDE_COLOR)
while True:
user_input = input(f"\n{USER_COLOR}You: {Style.RESET_ALL}")
if user_input.lower() == 'exit':
print_colored("Thank you for chatting. Goodbye!", CLAUDE_COLOR)
break
image_path = user_input.replace("'", "")
# Check if the input is a file path (potentially dragged and dropped)
if os.path.isfile(image_path):
user_input = input(f"{USER_COLOR}You (prompt for image): {Style.RESET_ALL}")
response = chat_with_claude(user_input, image_path)
else:
response = chat_with_claude(user_input)
# Check if the response contains code and format it
if "```" in response:
parts = response.split("```")
for i, part in enumerate(parts):
if i % 2 == 0:
print_colored(part, CLAUDE_COLOR)
else:
lines = part.split('\n')
language = lines[0].strip() if lines else ""
code = '\n'.join(lines[1:]) if len(lines) > 1 else ""
if language and code:
print_code(code, language)
elif code:
# If no language is specified but there is code, print it as plain text
print_colored(f"Code:\n{code}", CLAUDE_COLOR)
else:
# If there's no code (empty block), just print the part as is
print_colored(part, CLAUDE_COLOR)
if __name__ == "__main__":
main()