forked from docker/genai-stack
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
54 lines (45 loc) · 1.85 KB
/
utils.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
class BaseLogger:
def __init__(self) -> None:
self.info = print
def extract_title_and_question(input_string):
lines = input_string.strip().split("\n")
title = ""
question = ""
is_question = False # flag to know if we are inside a "Question" block
for line in lines:
if line.startswith("Title:"):
title = line.split("Title: ", 1)[1].strip()
elif line.startswith("Question:"):
question = line.split("Question: ", 1)[1].strip()
is_question = (
True # set the flag to True once we encounter a "Question:" line
)
elif is_question:
# if the line does not start with "Question:" but we are inside a "Question" block,
# then it is a continuation of the question
question += "\n" + line.strip()
return title, question
def create_vector_index(driver, dimension: int) -> None:
index_query = "CALL db.index.vector.createNodeIndex('stackoverflow', 'Question', 'embedding', $dimension, 'cosine')"
try:
driver.query(index_query, {"dimension": dimension})
except: # Already exists
pass
index_query = "CALL db.index.vector.createNodeIndex('top_answers', 'Answer', 'embedding', $dimension, 'cosine')"
try:
driver.query(index_query, {"dimension": dimension})
except: # Already exists
pass
def create_constraints(driver):
driver.query(
"CREATE CONSTRAINT question_id IF NOT EXISTS FOR (q:Question) REQUIRE (q.id) IS UNIQUE"
)
driver.query(
"CREATE CONSTRAINT answer_id IF NOT EXISTS FOR (a:Answer) REQUIRE (a.id) IS UNIQUE"
)
driver.query(
"CREATE CONSTRAINT user_id IF NOT EXISTS FOR (u:User) REQUIRE (u.id) IS UNIQUE"
)
driver.query(
"CREATE CONSTRAINT tag_name IF NOT EXISTS FOR (t:Tag) REQUIRE (t.name) IS UNIQUE"
)