ShynBui commited on
Commit
21705f0
·
verified ·
1 Parent(s): cfcdce8

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +24 -0
utils.py CHANGED
@@ -21,6 +21,30 @@ import ast
21
  import re
22
 
23
  parser = JsonOutputParser()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
  def query_as_list(db, query):
26
  res = db.run(query)
 
21
  import re
22
 
23
  parser = JsonOutputParser()
24
+ llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=os.environ['API_KEY'])
25
+ example_selector = SemanticSimilarityExampleSelector.from_examples(
26
+ examples,
27
+ OpenAIEmbeddings(model="text-embedding-3-small", api_key=os.environ['API_KEY']),
28
+ Chroma(persist_directory="data"),
29
+ # Chroma,
30
+ k=5,
31
+ input_keys=["input"],
32
+ )
33
+
34
+ db = SQLDatabase.from_uri("sqlite:///attendance_system.db")
35
+
36
+ employee = query_as_list(db, "SELECT FullName FROM Employee")
37
+
38
+ vector_db = Chroma.from_texts(employee, OpenAIEmbeddings(model="text-embedding-3-small", api_key=os.environ['API_KEY']))
39
+ retriever = vector_db.as_retriever(search_kwargs={"k": 15})
40
+ description = """Use to look up values to filter on. Input is an approximate spelling of the proper noun, output is \
41
+ valid proper nouns. Use the noun most similar to the search."""
42
+ retriever_tool = create_retriever_tool(
43
+ retriever,
44
+ name="search_proper_nouns",
45
+ description=description,
46
+ )
47
+
48
 
49
  def query_as_list(db, query):
50
  res = db.run(query)