Jacqkues commited on
Commit
0e0211f
·
verified ·
1 Parent(s): 51e5002

Delete agent.py

Browse files
Files changed (1) hide show
  1. agent.py +0 -62
agent.py DELETED
@@ -1,62 +0,0 @@
1
- from services.ai import rank_tables,generate_sql,generate_answer , correct_sql , evaluate_difficulty
2
- from services.utils import filter_tables
3
- from openai import OpenAI
4
- from database import Database
5
- from filesource import FileSource
6
- import os
7
- MAX_TABLE = 3
8
-
9
- client = OpenAI(
10
- base_url=os.getenv("LLM_ENDPOINT"),
11
- api_key=os.getenv("LLM_KEY")
12
- )
13
-
14
- def run_agent(database,prompt):
15
-
16
- retry = 5
17
- tables = database.get_tables_array()
18
-
19
- use_thinking = False
20
-
21
-
22
- if len(tables) > MAX_TABLE:
23
- print(f"using reranking because number of tables is greater than {MAX_TABLE}")
24
- ranked = rank_tables(prompt,tables)
25
- tables = filter_tables(0,ranked)[:MAX_TABLE]
26
-
27
- dif = int(evaluate_difficulty(client,prompt))
28
- if dif > 7:
29
- print("difficulty is > 7 so we enable thinking mode")
30
- use_thinking = True
31
- sql = generate_sql(client,prompt,tables,use_thinking)
32
- nb_try = 0
33
- success = False
34
- while nb_try < retry and not success:
35
- nb_try = nb_try + 1
36
- try:
37
- print("try to launch sql request")
38
- result = database.query(sql)
39
- success = True
40
- except Exception as e:
41
- print(f"Error : {e}")
42
- print("Try to self correct...")
43
- error = f"{type(e).__name__} - {str(e)}"
44
- if nb_try < retry - 2:
45
- sql = correct_sql(client,prompt,sql,tables,error,True)
46
- else:
47
- sql = correct_sql(client,prompt,sql,tables,error,False)
48
-
49
- print(sql)
50
-
51
- if success:
52
- print(result.to_markdown())
53
- return generate_answer(client,sql,prompt,result.to_markdown(),use_thinking)
54
-
55
-
56
- # db = Database("mysql://user:password@localhost:3306/Pokemon")
57
- # db.connect()
58
- # file = FileSource("./Wines.csv")
59
- # file.connect()
60
- # print(run_agent(file,"What is the quality og the win with the less of alcohol ?"))
61
-
62
-