Dhanushkumar commited on
Commit
c414791
·
verified ·
1 Parent(s): 18f73c8

Delete agent.py

Browse files
Files changed (1) hide show
  1. agent.py +0 -803
agent.py DELETED
@@ -1,803 +0,0 @@
1
- import os
2
- from dotenv import load_dotenv
3
- from typing import List, Dict, Any, Optional
4
- import tempfile
5
- import re
6
- import json
7
- import requests
8
- from urllib.parse import urlparse
9
- import pytesseract
10
- from PIL import Image, ImageDraw, ImageFont, ImageEnhance, ImageFilter
11
- import cmath
12
- import pandas as pd
13
- import uuid
14
- import numpy as np
15
- from code_interpreter import CodeInterpreter
16
-
17
- interpreter_instance = CodeInterpreter()
18
-
19
- from image_processing import *
20
-
21
- """Langraph"""
22
- from langgraph.graph import START, StateGraph, MessagesState
23
- from langchain_community.tools.tavily_search import TavilySearchResults
24
- from langchain_community.document_loaders import WikipediaLoader
25
- from langchain_community.document_loaders import ArxivLoader
26
- from langgraph.prebuilt import ToolNode, tools_condition
27
- from langchain_google_genai import ChatGoogleGenerativeAI
28
- from langchain_groq import ChatGroq
29
- from langchain_huggingface import (
30
- ChatHuggingFace,
31
- HuggingFaceEndpoint,
32
- HuggingFaceEmbeddings,
33
- )
34
- from langchain_community.vectorstores import SupabaseVectorStore
35
- from langchain_core.messages import SystemMessage, HumanMessage
36
- from langchain_core.tools import tool
37
- from langchain.tools.retriever import create_retriever_tool
38
- from supabase.client import Client, create_client
39
-
40
- load_dotenv()
41
-
42
- ### =============== BROWSER TOOLS =============== ###
43
- os.environ["TAVILY_API_KEY"] = "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"
44
- os.environ["GOOGLE_API_KEY"] = "AIzaSyCA47F1nv1sF9_y5IYzu_nZM325ZnIXfKU"
45
- os.environ["GROQ_API_KEY"] = "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"
46
- @tool
47
- def wiki_search(query: str) -> str:
48
- """Search Wikipedia for a query and return maximum 2 results.
49
-
50
- Args:
51
- query: The search query."""
52
- search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
53
- formatted_search_docs = "\n\n---\n\n".join(
54
- [
55
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
56
- for doc in search_docs
57
- ]
58
- )
59
- return {"wiki_results": formatted_search_docs}
60
-
61
-
62
- @tool
63
- def web_search(query: str) -> str:
64
- """Search Tavily for a query and return maximum 3 results.
65
-
66
- Args:
67
- query: The search query."""
68
- search_docs = TavilySearchResults(max_results=3).invoke(query=query)
69
- formatted_search_docs = "\n\n---\n\n".join(
70
- [
71
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
72
- for doc in search_docs
73
- ]
74
- )
75
- return {"web_results": formatted_search_docs}
76
-
77
-
78
- @tool
79
- def arxiv_search(query: str) -> str:
80
- """Search Arxiv for a query and return maximum 3 result.
81
-
82
- Args:
83
- query: The search query."""
84
- search_docs = ArxivLoader(query=query, load_max_docs=3).load()
85
- formatted_search_docs = "\n\n---\n\n".join(
86
- [
87
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
88
- for doc in search_docs
89
- ]
90
- )
91
- return {"arxiv_results": formatted_search_docs}
92
-
93
-
94
- ### =============== CODE INTERPRETER TOOLS =============== ###
95
-
96
-
97
- @tool
98
- def execute_code_multilang(code: str, language: str = "python") -> str:
99
- """Execute code in multiple languages (Python, Bash, SQL, C, Java) and return results.
100
-
101
- Args:
102
- code (str): The source code to execute.
103
- language (str): The language of the code. Supported: "python", "bash", "sql", "c", "java".
104
-
105
- Returns:
106
- A string summarizing the execution results (stdout, stderr, errors, plots, dataframes if any).
107
- """
108
- supported_languages = ["python", "bash", "sql", "c", "java"]
109
- language = language.lower()
110
-
111
- if language not in supported_languages:
112
- return f"❌ Unsupported language: {language}. Supported languages are: {', '.join(supported_languages)}"
113
-
114
- result = interpreter_instance.execute_code(code, language=language)
115
-
116
- response = []
117
-
118
- if result["status"] == "success":
119
- response.append(f"✅ Code executed successfully in **{language.upper()}**")
120
-
121
- if result.get("stdout"):
122
- response.append(
123
- "\n**Standard Output:**\n```\n" + result["stdout"].strip() + "\n```"
124
- )
125
-
126
- if result.get("stderr"):
127
- response.append(
128
- "\n**Standard Error (if any):**\n```\n"
129
- + result["stderr"].strip()
130
- + "\n```"
131
- )
132
-
133
- if result.get("result") is not None:
134
- response.append(
135
- "\n**Execution Result:**\n```\n"
136
- + str(result["result"]).strip()
137
- + "\n```"
138
- )
139
-
140
- if result.get("dataframes"):
141
- for df_info in result["dataframes"]:
142
- response.append(
143
- f"\n**DataFrame `{df_info['name']}` (Shape: {df_info['shape']})**"
144
- )
145
- df_preview = pd.DataFrame(df_info["head"])
146
- response.append("First 5 rows:\n```\n" + str(df_preview) + "\n```")
147
-
148
- if result.get("plots"):
149
- response.append(
150
- f"\n**Generated {len(result['plots'])} plot(s)** (Image data returned separately)"
151
- )
152
-
153
- else:
154
- response.append(f"❌ Code execution failed in **{language.upper()}**")
155
- if result.get("stderr"):
156
- response.append(
157
- "\n**Error Log:**\n```\n" + result["stderr"].strip() + "\n```"
158
- )
159
-
160
- return "\n".join(response)
161
-
162
-
163
- ### =============== MATHEMATICAL TOOLS =============== ###
164
-
165
-
166
- @tool
167
- def multiply(a: float, b: float) -> float:
168
- """
169
- Multiplies two numbers.
170
-
171
- Args:
172
- a (float): the first number
173
- b (float): the second number
174
- """
175
- return a * b
176
-
177
-
178
- @tool
179
- def add(a: float, b: float) -> float:
180
- """
181
- Adds two numbers.
182
-
183
- Args:
184
- a (float): the first number
185
- b (float): the second number
186
- """
187
- return a + b
188
-
189
-
190
- @tool
191
- def subtract(a: float, b: float) -> int:
192
- """
193
- Subtracts two numbers.
194
-
195
- Args:
196
- a (float): the first number
197
- b (float): the second number
198
- """
199
- return a - b
200
-
201
-
202
- @tool
203
- def divide(a: float, b: float) -> float:
204
- """
205
- Divides two numbers.
206
-
207
- Args:
208
- a (float): the first float number
209
- b (float): the second float number
210
- """
211
- if b == 0:
212
- raise ValueError("Cannot divided by zero.")
213
- return a / b
214
-
215
-
216
- @tool
217
- def modulus(a: int, b: int) -> int:
218
- """
219
- Get the modulus of two numbers.
220
-
221
- Args:
222
- a (int): the first number
223
- b (int): the second number
224
- """
225
- return a % b
226
-
227
-
228
- @tool
229
- def power(a: float, b: float) -> float:
230
- """
231
- Get the power of two numbers.
232
-
233
- Args:
234
- a (float): the first number
235
- b (float): the second number
236
- """
237
- return a**b
238
-
239
-
240
- @tool
241
- def square_root(a: float) -> float | complex:
242
- """
243
- Get the square root of a number.
244
-
245
- Args:
246
- a (float): the number to get the square root of
247
- """
248
- if a >= 0:
249
- return a**0.5
250
- return cmath.sqrt(a)
251
-
252
-
253
- ### =============== DOCUMENT PROCESSING TOOLS =============== ###
254
-
255
-
256
- @tool
257
- def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
258
- """
259
- Save content to a file and return the path.
260
-
261
- Args:
262
- content (str): the content to save to the file
263
- filename (str, optional): the name of the file. If not provided, a random name file will be created.
264
- """
265
- temp_dir = tempfile.gettempdir()
266
- if filename is None:
267
- temp_file = tempfile.NamedTemporaryFile(delete=False, dir=temp_dir)
268
- filepath = temp_file.name
269
- else:
270
- filepath = os.path.join(temp_dir, filename)
271
-
272
- with open(filepath, "w") as f:
273
- f.write(content)
274
-
275
- return f"File saved to {filepath}. You can read this file to process its contents."
276
-
277
-
278
- @tool
279
- def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
280
- """
281
- Download a file from a URL and save it to a temporary location.
282
-
283
- Args:
284
- url (str): the URL of the file to download.
285
- filename (str, optional): the name of the file. If not provided, a random name file will be created.
286
- """
287
- try:
288
- # Parse URL to get filename if not provided
289
- if not filename:
290
- path = urlparse(url).path
291
- filename = os.path.basename(path)
292
- if not filename:
293
- filename = f"downloaded_{uuid.uuid4().hex[:8]}"
294
-
295
- # Create temporary file
296
- temp_dir = tempfile.gettempdir()
297
- filepath = os.path.join(temp_dir, filename)
298
-
299
- # Download the file
300
- response = requests.get(url, stream=True)
301
- response.raise_for_status()
302
-
303
- # Save the file
304
- with open(filepath, "wb") as f:
305
- for chunk in response.iter_content(chunk_size=8192):
306
- f.write(chunk)
307
-
308
- return f"File downloaded to {filepath}. You can read this file to process its contents."
309
- except Exception as e:
310
- return f"Error downloading file: {str(e)}"
311
-
312
-
313
- @tool
314
- def extract_text_from_image(image_path: str) -> str:
315
- """
316
- Extract text from an image using OCR library pytesseract (if available).
317
-
318
- Args:
319
- image_path (str): the path to the image file.
320
- """
321
- try:
322
- # Open the image
323
- image = Image.open(image_path)
324
-
325
- # Extract text from the image
326
- text = pytesseract.image_to_string(image)
327
-
328
- return f"Extracted text from image:\n\n{text}"
329
- except Exception as e:
330
- return f"Error extracting text from image: {str(e)}"
331
-
332
-
333
- @tool
334
- def analyze_csv_file(file_path: str, query: str) -> str:
335
- """
336
- Analyze a CSV file using pandas and answer a question about it.
337
-
338
- Args:
339
- file_path (str): the path to the CSV file.
340
- query (str): Question about the data
341
- """
342
- try:
343
- # Read the CSV file
344
- df = pd.read_csv(file_path)
345
-
346
- # Run various analyses based on the query
347
- result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
348
- result += f"Columns: {', '.join(df.columns)}\n\n"
349
-
350
- # Add summary statistics
351
- result += "Summary statistics:\n"
352
- result += str(df.describe())
353
-
354
- return result
355
-
356
- except Exception as e:
357
- return f"Error analyzing CSV file: {str(e)}"
358
-
359
-
360
- @tool
361
- def analyze_excel_file(file_path: str, query: str) -> str:
362
- """
363
- Analyze an Excel file using pandas and answer a question about it.
364
-
365
- Args:
366
- file_path (str): the path to the Excel file.
367
- query (str): Question about the data
368
- """
369
- try:
370
- # Read the Excel file
371
- df = pd.read_excel(file_path)
372
-
373
- # Run various analyses based on the query
374
- result = (
375
- f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
376
- )
377
- result += f"Columns: {', '.join(df.columns)}\n\n"
378
-
379
- # Add summary statistics
380
- result += "Summary statistics:\n"
381
- result += str(df.describe())
382
-
383
- return result
384
-
385
- except Exception as e:
386
- return f"Error analyzing Excel file: {str(e)}"
387
-
388
-
389
- ### ============== IMAGE PROCESSING AND GENERATION TOOLS =============== ###
390
-
391
-
392
- @tool
393
- def analyze_image(image_base64: str) -> Dict[str, Any]:
394
- """
395
- Analyze basic properties of an image (size, mode, color analysis, thumbnail preview).
396
-
397
- Args:
398
- image_base64 (str): Base64 encoded image string
399
-
400
- Returns:
401
- Dictionary with analysis result
402
- """
403
- try:
404
- img = decode_image(image_base64)
405
- width, height = img.size
406
- mode = img.mode
407
-
408
- if mode in ("RGB", "RGBA"):
409
- arr = np.array(img)
410
- avg_colors = arr.mean(axis=(0, 1))
411
- dominant = ["Red", "Green", "Blue"][np.argmax(avg_colors[:3])]
412
- brightness = avg_colors.mean()
413
- color_analysis = {
414
- "average_rgb": avg_colors.tolist(),
415
- "brightness": brightness,
416
- "dominant_color": dominant,
417
- }
418
- else:
419
- color_analysis = {"note": f"No color analysis for mode {mode}"}
420
-
421
- thumbnail = img.copy()
422
- thumbnail.thumbnail((100, 100))
423
- thumb_path = save_image(thumbnail, "thumbnails")
424
- thumbnail_base64 = encode_image(thumb_path)
425
-
426
- return {
427
- "dimensions": (width, height),
428
- "mode": mode,
429
- "color_analysis": color_analysis,
430
- "thumbnail": thumbnail_base64,
431
- }
432
- except Exception as e:
433
- return {"error": str(e)}
434
-
435
-
436
- @tool
437
- def transform_image(
438
- image_base64: str, operation: str, params: Optional[Dict[str, Any]] = None
439
- ) -> Dict[str, Any]:
440
- """
441
- Apply transformations: resize, rotate, crop, flip, brightness, contrast, blur, sharpen, grayscale.
442
-
443
- Args:
444
- image_base64 (str): Base64 encoded input image
445
- operation (str): Transformation operation
446
- params (Dict[str, Any], optional): Parameters for the operation
447
-
448
- Returns:
449
- Dictionary with transformed image (base64)
450
- """
451
- try:
452
- img = decode_image(image_base64)
453
- params = params or {}
454
-
455
- if operation == "resize":
456
- img = img.resize(
457
- (
458
- params.get("width", img.width // 2),
459
- params.get("height", img.height // 2),
460
- )
461
- )
462
- elif operation == "rotate":
463
- img = img.rotate(params.get("angle", 90), expand=True)
464
- elif operation == "crop":
465
- img = img.crop(
466
- (
467
- params.get("left", 0),
468
- params.get("top", 0),
469
- params.get("right", img.width),
470
- params.get("bottom", img.height),
471
- )
472
- )
473
- elif operation == "flip":
474
- if params.get("direction", "horizontal") == "horizontal":
475
- img = img.transpose(Image.FLIP_LEFT_RIGHT)
476
- else:
477
- img = img.transpose(Image.FLIP_TOP_BOTTOM)
478
- elif operation == "adjust_brightness":
479
- img = ImageEnhance.Brightness(img).enhance(params.get("factor", 1.5))
480
- elif operation == "adjust_contrast":
481
- img = ImageEnhance.Contrast(img).enhance(params.get("factor", 1.5))
482
- elif operation == "blur":
483
- img = img.filter(ImageFilter.GaussianBlur(params.get("radius", 2)))
484
- elif operation == "sharpen":
485
- img = img.filter(ImageFilter.SHARPEN)
486
- elif operation == "grayscale":
487
- img = img.convert("L")
488
- else:
489
- return {"error": f"Unknown operation: {operation}"}
490
-
491
- result_path = save_image(img)
492
- result_base64 = encode_image(result_path)
493
- return {"transformed_image": result_base64}
494
-
495
- except Exception as e:
496
- return {"error": str(e)}
497
-
498
-
499
- @tool
500
- def draw_on_image(
501
- image_base64: str, drawing_type: str, params: Dict[str, Any]
502
- ) -> Dict[str, Any]:
503
- """
504
- Draw shapes (rectangle, circle, line) or text onto an image.
505
-
506
- Args:
507
- image_base64 (str): Base64 encoded input image
508
- drawing_type (str): Drawing type
509
- params (Dict[str, Any]): Drawing parameters
510
-
511
- Returns:
512
- Dictionary with result image (base64)
513
- """
514
- try:
515
- img = decode_image(image_base64)
516
- draw = ImageDraw.Draw(img)
517
- color = params.get("color", "red")
518
-
519
- if drawing_type == "rectangle":
520
- draw.rectangle(
521
- [params["left"], params["top"], params["right"], params["bottom"]],
522
- outline=color,
523
- width=params.get("width", 2),
524
- )
525
- elif drawing_type == "circle":
526
- x, y, r = params["x"], params["y"], params["radius"]
527
- draw.ellipse(
528
- (x - r, y - r, x + r, y + r),
529
- outline=color,
530
- width=params.get("width", 2),
531
- )
532
- elif drawing_type == "line":
533
- draw.line(
534
- (
535
- params["start_x"],
536
- params["start_y"],
537
- params["end_x"],
538
- params["end_y"],
539
- ),
540
- fill=color,
541
- width=params.get("width", 2),
542
- )
543
- elif drawing_type == "text":
544
- font_size = params.get("font_size", 20)
545
- try:
546
- font = ImageFont.truetype("arial.ttf", font_size)
547
- except IOError:
548
- font = ImageFont.load_default()
549
- draw.text(
550
- (params["x"], params["y"]),
551
- params.get("text", "Text"),
552
- fill=color,
553
- font=font,
554
- )
555
- else:
556
- return {"error": f"Unknown drawing type: {drawing_type}"}
557
-
558
- result_path = save_image(img)
559
- result_base64 = encode_image(result_path)
560
- return {"result_image": result_base64}
561
-
562
- except Exception as e:
563
- return {"error": str(e)}
564
-
565
-
566
- @tool
567
- def generate_simple_image(
568
- image_type: str,
569
- width: int = 500,
570
- height: int = 500,
571
- params: Optional[Dict[str, Any]] = None,
572
- ) -> Dict[str, Any]:
573
- """
574
- Generate a simple image (gradient, noise, pattern, chart).
575
-
576
- Args:
577
- image_type (str): Type of image
578
- width (int), height (int)
579
- params (Dict[str, Any], optional): Specific parameters
580
-
581
- Returns:
582
- Dictionary with generated image (base64)
583
- """
584
- try:
585
- params = params or {}
586
-
587
- if image_type == "gradient":
588
- direction = params.get("direction", "horizontal")
589
- start_color = params.get("start_color", (255, 0, 0))
590
- end_color = params.get("end_color", (0, 0, 255))
591
-
592
- img = Image.new("RGB", (width, height))
593
- draw = ImageDraw.Draw(img)
594
-
595
- if direction == "horizontal":
596
- for x in range(width):
597
- r = int(
598
- start_color[0] + (end_color[0] - start_color[0]) * x / width
599
- )
600
- g = int(
601
- start_color[1] + (end_color[1] - start_color[1]) * x / width
602
- )
603
- b = int(
604
- start_color[2] + (end_color[2] - start_color[2]) * x / width
605
- )
606
- draw.line([(x, 0), (x, height)], fill=(r, g, b))
607
- else:
608
- for y in range(height):
609
- r = int(
610
- start_color[0] + (end_color[0] - start_color[0]) * y / height
611
- )
612
- g = int(
613
- start_color[1] + (end_color[1] - start_color[1]) * y / height
614
- )
615
- b = int(
616
- start_color[2] + (end_color[2] - start_color[2]) * y / height
617
- )
618
- draw.line([(0, y), (width, y)], fill=(r, g, b))
619
-
620
- elif image_type == "noise":
621
- noise_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
622
- img = Image.fromarray(noise_array, "RGB")
623
-
624
- else:
625
- return {"error": f"Unsupported image_type {image_type}"}
626
-
627
- result_path = save_image(img)
628
- result_base64 = encode_image(result_path)
629
- return {"generated_image": result_base64}
630
-
631
- except Exception as e:
632
- return {"error": str(e)}
633
-
634
-
635
- @tool
636
- def combine_images(
637
- images_base64: List[str], operation: str, params: Optional[Dict[str, Any]] = None
638
- ) -> Dict[str, Any]:
639
- """
640
- Combine multiple images (collage, stack, blend).
641
-
642
- Args:
643
- images_base64 (List[str]): List of base64 images
644
- operation (str): Combination type
645
- params (Dict[str, Any], optional)
646
-
647
- Returns:
648
- Dictionary with combined image (base64)
649
- """
650
- try:
651
- images = [decode_image(b64) for b64 in images_base64]
652
- params = params or {}
653
-
654
- if operation == "stack":
655
- direction = params.get("direction", "horizontal")
656
- if direction == "horizontal":
657
- total_width = sum(img.width for img in images)
658
- max_height = max(img.height for img in images)
659
- new_img = Image.new("RGB", (total_width, max_height))
660
- x = 0
661
- for img in images:
662
- new_img.paste(img, (x, 0))
663
- x += img.width
664
- else:
665
- max_width = max(img.width for img in images)
666
- total_height = sum(img.height for img in images)
667
- new_img = Image.new("RGB", (max_width, total_height))
668
- y = 0
669
- for img in images:
670
- new_img.paste(img, (0, y))
671
- y += img.height
672
- else:
673
- return {"error": f"Unsupported combination operation {operation}"}
674
-
675
- result_path = save_image(new_img)
676
- result_base64 = encode_image(result_path)
677
- return {"combined_image": result_base64}
678
-
679
- except Exception as e:
680
- return {"error": str(e)}
681
-
682
-
683
- # load the system prompt from the file
684
- with open("system_prompt.txt", "r", encoding="utf-8") as f:
685
- system_prompt = f.read()
686
- print(system_prompt)
687
-
688
- # System message
689
- sys_msg = SystemMessage(content=system_prompt)
690
-
691
- # build a retriever
692
- embeddings = HuggingFaceEmbeddings(
693
- model_name="sentence-transformers/all-mpnet-base-v2"
694
- ) # dim=768
695
- supabase_url = "https://kyeropjbpzxfcypbxqoi.supabase.co"
696
- supabase_key = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6Imt5ZXJvcGpicHp4ZmN5cGJ4cW9pIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NDg0ODYzNDEsImV4cCI6MjA2NDA2MjM0MX0._7ct7b36PILtOs4ycgPk1Tgk6dKNp2bbTjB-UBA3uGg"
697
- supabase: Client = create_client(supabase_url, supabase_key)
698
-
699
- vector_store = SupabaseVectorStore(
700
- client=supabase,
701
- embedding=embeddings,
702
- table_name="documents2",
703
- query_name="match_documents_2",
704
- )
705
- create_retriever_tool = create_retriever_tool(
706
- retriever=vector_store.as_retriever(),
707
- name="Question Search",
708
- description="A tool to retrieve similar questions from a vector store.",
709
- )
710
-
711
-
712
- tools = [
713
- web_search,
714
- wiki_search,
715
- arxiv_search,
716
- multiply,
717
- add,
718
- subtract,
719
- divide,
720
- modulus,
721
- power,
722
- square_root,
723
- save_and_read_file,
724
- download_file_from_url,
725
- extract_text_from_image,
726
- analyze_csv_file,
727
- analyze_excel_file,
728
- execute_code_multilang,
729
- analyze_image,
730
- transform_image,
731
- draw_on_image,
732
- generate_simple_image,
733
- combine_images,
734
- ]
735
-
736
-
737
- # Build graph function
738
- def build_graph(provider: str = "groq"):
739
- """Build the graph"""
740
- # Load environment variables from .env file
741
- if provider == "groq":
742
- # Groq https://console.groq.com/docs/models
743
- llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
744
- elif provider == "huggingface":
745
- # TODO: Add huggingface endpoint
746
- llm = ChatHuggingFace(
747
- llm=HuggingFaceEndpoint(
748
- repo_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
749
- task="text-generation", # for chat‐style use “text-generation”
750
- max_new_tokens=1024,
751
- do_sample=False,
752
- repetition_penalty=1.03,
753
- temperature=0,
754
- ),
755
- verbose=True,
756
- )
757
- else:
758
- raise ValueError("Invalid provider. Choose 'groq' or 'huggingface'.")
759
- # Bind tools to LLM
760
- llm_with_tools = llm.bind_tools(tools)
761
-
762
- # Node
763
- def assistant(state: MessagesState):
764
- """Assistant node"""
765
- return {"messages": [llm_with_tools.invoke(state["messages"])]}
766
-
767
- def retriever(state: MessagesState):
768
- """Retriever node"""
769
- similar_question = vector_store.similarity_search(state["messages"][0].content)
770
-
771
- if similar_question: # Check if the list is not empty
772
- example_msg = HumanMessage(
773
- content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
774
- )
775
- return {"messages": [sys_msg] + state["messages"] + [example_msg]}
776
- else:
777
- # Handle the case when no similar questions are found
778
- return {"messages": [sys_msg] + state["messages"]}
779
-
780
- builder = StateGraph(MessagesState)
781
- builder.add_node("retriever", retriever)
782
- builder.add_node("assistant", assistant)
783
- builder.add_node("tools", ToolNode(tools))
784
- builder.add_edge(START, "retriever")
785
- builder.add_edge("retriever", "assistant")
786
- builder.add_conditional_edges(
787
- "assistant",
788
- tools_condition,
789
- )
790
- builder.add_edge("tools", "assistant")
791
-
792
- # Compile graph
793
- return builder.compile()
794
-
795
-
796
- # test
797
- if __name__ == "__main__":
798
- question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
799
- graph = build_graph(provider="groq")
800
- messages = [HumanMessage(content=question)]
801
- messages = graph.invoke({"messages": messages})
802
- for m in messages["messages"]:
803
- m.pretty_print()