dkdaniz commited on
Commit
2c3245b
1 Parent(s): 751c042

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +44 -0
main.py CHANGED
@@ -94,3 +94,47 @@ async def predict(data: Predict):
94
  return prompt_response_dict
95
  else:
96
  raise HTTPException(status_code=400, detail="Prompt Incorrect")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
94
  return prompt_response_dict
95
  else:
96
  raise HTTPException(status_code=400, detail="Prompt Incorrect")
97
+
98
+ @app.get("/run_ingest")
99
+ def run_ingest_route():
100
+ try:
101
+ if os.path.exists(PERSIST_DIRECTORY):
102
+ try:
103
+ shutil.rmtree(PERSIST_DIRECTORY)
104
+ except OSError as e:
105
+ print(f"Error: {e.filename} - {e.strerror}.")
106
+ else:
107
+ print("The directory does not exist")
108
+
109
+ run_langest_commands = ["python", "ingest.py"]
110
+ if DEVICE_TYPE == "cpu":
111
+ run_langest_commands.append("--device_type")
112
+ run_langest_commands.append(DEVICE_TYPE)
113
+
114
+ result = subprocess.run(run_langest_commands, capture_output=True)
115
+
116
+ if result.returncode != 0:
117
+ raise HTTPException(status_code=400, detail="Script execution failed: {}")
118
+
119
+ # load the vectorstore
120
+ DB = Chroma(
121
+ persist_directory=PERSIST_DIRECTORY,
122
+ embedding_function=EMBEDDINGS,
123
+ client_settings=CHROMA_SETTINGS,
124
+ )
125
+ RETRIEVER = DB.as_retriever()
126
+ prompt, memory = get_prompt_template(promptTemplate_type="llama", history=False)
127
+
128
+ QA = RetrievalQA.from_chain_type(
129
+ llm=LLM,
130
+ chain_type="stuff",
131
+ retriever=RETRIEVER,
132
+ return_source_documents=SHOW_SOURCES,
133
+ chain_type_kwargs={
134
+ "prompt": prompt,
135
+ },
136
+ )
137
+ return "Script executed successfully: {}".format(result.stdout.decode("utf-8"))
138
+ except Exception as e:
139
+ raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
140
+