anubhav77 commited on
Commit
8d7feb0
·
1 Parent(s): ee4103c
Files changed (2) hide show
  1. src/chromaIntf.py +15 -5
  2. src/main.py +4 -4
src/chromaIntf.py CHANGED
@@ -9,6 +9,8 @@ from llm.llmFactory import LLMFactory
9
  from datetime import datetime
10
  import baseInfra.dropbox_handler as dbh
11
  from baseInfra.dbInterface import DbInterface
 
 
12
 
13
  class ChromaIntf():
14
  def __init__(self):
@@ -96,7 +98,7 @@ class ChromaIntf():
96
  return retVal
97
 
98
 
99
- def addText(self,inStr:str,metadata):
100
  metadata=metadata.dict()
101
  if "timestamp" not in metadata.keys():
102
  metadata['timestamp']=datetime.now().isoformat()
@@ -117,17 +119,25 @@ class ChromaIntf():
117
  docs = [
118
  Document(page_content=inStr, metadata=metadata)]
119
  try:
120
- return self.vectorstore.add_documents(docs,ids=[metadata.ID])
121
  except:
122
  print("inside expect of addText")
123
- return self.vectorstore.add_documents(docs,ids=[metadata['ID']])
124
 
125
  async def listDocs(self):
126
- return self.vectorstore._client.get_collection(self.vectorstore._LANGCHAIN_DEFAULT_COLLECTION_NAME,embedding_function=self.embedding)
 
127
 
128
 
129
  async def persist(self):
130
  await self.vectorstore.persist()
131
- await dbh.backupFolder("db")
 
 
 
 
 
 
 
132
 
133
 
 
9
  from datetime import datetime
10
  import baseInfra.dropbox_handler as dbh
11
  from baseInfra.dbInterface import DbInterface
12
+ from uuid import UUID
13
+
14
 
15
  class ChromaIntf():
16
  def __init__(self):
 
98
  return retVal
99
 
100
 
101
+ async def addText(self,inStr:str,metadata):
102
  metadata=metadata.dict()
103
  if "timestamp" not in metadata.keys():
104
  metadata['timestamp']=datetime.now().isoformat()
 
119
  docs = [
120
  Document(page_content=inStr, metadata=metadata)]
121
  try:
122
+ return await self.vectorstore.add_documents(docs,ids=[metadata.ID])
123
  except:
124
  print("inside expect of addText")
125
+ return await self.vectorstore.add_documents(docs,ids=[metadata['ID']])
126
 
127
  async def listDocs(self):
128
+ collectionInfo=self.vectorstore._client.get_collection(self.vectorstore._LANGCHAIN_DEFAULT_COLLECTION_NAME,embedding_function=self.embedding)[0]
129
+ return self.vectorstore._client._get(collection_id=self._uuid(collectionInfo.id))
130
 
131
 
132
  async def persist(self):
133
  await self.vectorstore.persist()
134
+ return await dbh.backupFolder("db")
135
+
136
+ def _uuid(uuid_str: str) -> UUID:
137
+ try:
138
+ return UUID(uuid_str)
139
+ except ValueError:
140
+ print("Error generating uuid")
141
+ raise ValueError(f"Could not parse {uuid_str} as a UUID")
142
 
143
 
src/main.py CHANGED
@@ -36,8 +36,8 @@ async def catch_exceptions_middleware(
36
  request: Request, call_next: Callable[[Request], Any]
37
  ) -> Response:
38
  try:
39
- print("In exception cater middleware")
40
- print(request.headers)
41
  #print(await request.body())
42
  return await call_next(request)
43
  except Exception as e:
@@ -66,7 +66,7 @@ async def get_matching_docs(inStr: str ) -> Any:
66
  TODO: Add parameter for type of query and number of docs to return
67
  TODO: Add parameter to return the source information as well
68
  """
69
- return chromaIntf.getRelevantDocs(inStr)
70
 
71
  @app.post(api_base+"/addTextDocument")
72
  async def add_text_document(inDoc: DocWithMeta ) -> Any:
@@ -75,7 +75,7 @@ async def add_text_document(inDoc: DocWithMeta ) -> Any:
75
  """
76
  print("Received request for")
77
  print(inDoc)
78
- return chromaIntf.addText(inDoc.text,inDoc.metadata)
79
 
80
  @app.get(api_base+"/persist")
81
  async def persist_db():
 
36
  request: Request, call_next: Callable[[Request], Any]
37
  ) -> Response:
38
  try:
39
+ #print("In exception cater middleware")
40
+ #print(request.headers)
41
  #print(await request.body())
42
  return await call_next(request)
43
  except Exception as e:
 
66
  TODO: Add parameter for type of query and number of docs to return
67
  TODO: Add parameter to return the source information as well
68
  """
69
+ return await chromaIntf.getRelevantDocs(inStr)
70
 
71
  @app.post(api_base+"/addTextDocument")
72
  async def add_text_document(inDoc: DocWithMeta ) -> Any:
 
75
  """
76
  print("Received request for")
77
  print(inDoc)
78
+ return await chromaIntf.addText(inDoc.text,inDoc.metadata)
79
 
80
  @app.get(api_base+"/persist")
81
  async def persist_db():