add docs

service/api.py

@@ -15,10 +15,22 @@ class API():
     def __init__(self) -> None:
         pass
 
+    def search(self,query,**kwargs):
+        if not query:
+            return []
+        docs = vector_store.search(query,**kwargs)
+        return docs
+    
+    def marginal_search(self,query,**kwargs):
+        if not query:
+            return []
+        docs = vector_store.marginal_search(query,**kwargs)
+        return docs
+
     def get_suggestion(self,query):
         if not query:
             return []
-        docs = vector_store.marginal_search(query)
+        docs = vector_store.marginal_search(query,lambda_mult=0.6)
         res = []
         for doc in docs:
             name = doc.metadata.get('title_llm','')
@@ -28,7 +40,6 @@ class API():
             res.append({"name":name,"url":url})
         # logger.write_log({"query":query,"api":'auto_complete',"result":res})
         return res
-        
     
     def get_answer(self,query):
         docs = vector_store.search(query)

service/crawler.py

@@ -0,0 +1,100 @@
+import json
+import re
+from tqdm import tqdm
+from bs4 import BeautifulSoup
+from langchain.document_loaders.recursive_url_loader import RecursiveUrlLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from typing import List
+from service.llm import Bot
+from util.schema import MyDoc
+from typing import List
+
+class Crawler(RecursiveUrlLoader):
+
+    HEADERS = {
+        'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/56.0.2924.76 Safari/537.36'
+    }
+    CHUNK_SIZE = 2000
+    TEMPLATE = """
+        You are an expert in refining and summarization based on the title, description of the document to make it more concise and informative. Please generate a refined title, description, and keywords for the document in json format. 
+        Remember, the generated keywords should be less than 3 words, the generated title should be less than 10 words, and the generated description should be less than 100 words.
+        Here is a legal example of the json format:
+        {{
+            "title": "The Title of the Document",
+            "description": "The Description of the Document",
+            "keywords": ["keyword1", "keyword2", "keyword3"]
+        }}
+    """
+
+    def __init__(self,keywords=[]):
+        self.keywords = keywords
+        self.splitter = RecursiveCharacterTextSplitter(chunk_size=self.CHUNK_SIZE, chunk_overlap=200)
+        self.bot = Bot(template=self.TEMPLATE)
+        super().__init__(
+            url="",
+            extractor=self._extractor,
+            timeout=600,
+            headers=self.HEADERS,
+            exclude_dirs=[],
+            prevent_outside=True
+        )
+
+    def _extractor(self,html: str) -> str:
+        text = BeautifulSoup(html, "lxml").get_text(" ")
+        text = re.sub(r"\n\n+", "\n\n", text)
+        text = re.sub(r"\s\s+", "  ", text).strip()
+        if self.check_keywords(text):
+            return text
+        return
+    
+    def check_keywords(self,text):
+        if not self.keywords:
+            return True
+        for key in self.keywords:
+            if key in text:
+                return True
+        return False
+
+    def __call__(self,url,visited=set(),max_depth=2,prevent_outside=True):
+        self.url = url
+        self.max_depth = max_depth
+        self.prevent_outside = prevent_outside
+        docs = list(self._get_child_links_recursive(url, visited))
+        return docs
+    
+    def from_docs(self,docs):
+        return MyDoc.from_docs(docs)
+    
+    def chunk(self,docs):
+        res = []
+        for doc in docs:
+            res += self.splitter.split_documents([doc])
+        return res
+    
+    def llm_augment(self,docs:List[MyDoc]):
+        question_template = """
+        Please generate a refined title, description, and keywords for the document
+        1. title:{title}
+        2. desription:{description}
+        3. content:{content}
+        """     
+        res =[]
+        for doc in tqdm(docs):
+            question = question_template.format(
+                title=doc.title,
+                description=doc.description,
+                content=doc.page_content
+            )
+            response = self.bot.custom_call(question=question)
+            try:
+                response = response.strip('`').strip('json\n')
+                data = json.loads(response)
+                doc.keywords = data['keywords']
+                doc.description_llm = data['description']
+                doc.title_llm = data['title']
+                res.append(doc)
+            except Exception as e:
+                print(response)
+        return res   
+    
+crawler = Crawler()

service/vector_store.py

@@ -38,13 +38,57 @@ class VectorStore(Pinecone):
             pinecone.create_index(name=self.INDEX_NAME, metric="cosine", dimension=self.dims)
         Pinecone.from_documents(docs, self.embeddings, index_name=self.INDEX_NAME)
 
-    # @timing
-    def search(self,query):
-        return self.similarity_search(query)
+    def search(self,
+        query: str,
+        k: int = 4,
+        filter: Optional[dict] = None,
+        namespace: Optional[str] = None,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Return pinecone documents most similar to query.
 
-    # @timing
-    def marginal_search(self,query,k=5):
-        return self.max_marginal_relevance_search(query,k=k)
+        Args:
+            query: Text to look up documents similar to.
+            k: Number of Documents to return. Defaults to 4.
+            filter: Dictionary of argument(s) to filter on metadata
+            namespace: Namespace to search in. Default will search in '' namespace.
+
+        Returns:
+            List of Documents most similar to the query and score for each
+        """
+        docs_and_scores = self.similarity_search_with_score(
+            query, k=k, filter=filter, namespace=namespace, **kwargs
+        )
+        return [doc for doc, _ in docs_and_scores]
+
+    def marginal_search(self,
+        query: str,
+        k: int = 4,
+        fetch_k: int = 20,
+        lambda_mult: float = 0.5,
+        filter: Optional[dict] = None,
+        namespace: Optional[str] = None,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Return docs selected using the maximal marginal relevance.
+
+        Maximal marginal relevance optimizes for similarity to query AND diversity
+        among selected documents.
+
+        Args:
+            query: Text to look up documents similar to.
+            k: Number of Documents to return. Defaults to 4.
+            fetch_k: Number of Documents to fetch to pass to MMR algorithm.
+            lambda_mult: Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding
+                        to maximum diversity and 1 to minimum diversity.
+                        Defaults to 0.5.
+        Returns:
+            List of Documents selected by maximal marginal relevance.
+        """
+        embedding = self._embed_query(query)
+        return self.max_marginal_relevance_search_by_vector(
+            embedding, k, fetch_k, lambda_mult, filter, namespace
+        )
         
     def jsonfy(self,docs):
         docs = [doc.dict() for doc in docs]

util/schema.py

@@ -39,4 +39,82 @@ class Card:
             "title": self.title,
             "description": self.description,
             "url": self.url
-        }
+        }
+
+
+
+class MyDoc:
+
+    def __init__(self, content, metadata):
+        self._content = content
+        self._metadata = {k:v for k,v in metadata.items() if v is not None and v.strip()!= ''}
+    
+    @staticmethod
+    def from_docs(docs):
+        res = []
+        for doc in docs:
+            res.append(MyDoc(doc.page_content, doc.metadata))
+        return res
+    
+    @property
+    def idx(self):
+        return self.metadata.get('idx', -1)
+
+    @property
+    def page_content(self):
+        return self._content
+    
+    @property
+    def metadata(self):
+        return self._metadata
+    
+    @page_content.setter
+    def page_content(self, value):
+        self._page_content = value
+
+    @metadata.setter
+    def metadata(self, value):
+        self._metadata = value
+
+    @property
+    def title(self):
+        return self.metadata.get('title', '')
+    
+    @property
+    def source(self):
+        return self.metadata.get('source', '')
+    
+    @property
+    def description(self):
+        return self.metadata.get('description', '')
+    
+    @property
+    def language(self):
+        return self.metadata.get('language', '')
+    
+    @property
+    def keywords(self):
+        return self.metadata.get('keywords', [])
+    
+    @property
+    def description_llm(self):
+        return self.metadata.get('description_llm', '')
+    
+    @property
+    def title_llm(self):
+        return self.metadata.get('title_llm', '')
+    
+    @description_llm.setter
+    def description_llm(self, value):
+        self.metadata['description_llm'] = value
+    
+    @title_llm.setter
+    def title_llm(self, value):
+        self.metadata['title_llm'] = value
+    
+    @keywords.setter
+    def keywords(self, value):
+        self.metadata['keywords'] = value
+    
+    def __str__(self):
+        return f"Source: {self.source}"