Merge pull request #110 from ant-xuexiao/rag_pr

feat: enable to add knowledge form git file or issues

Author: RaoHai

server/.env.example

@@ -22,3 +22,5 @@ AUTH0_CLIENT_ID=auth0_client_id
 AUTH0_CLIENT_SECRET=auth0_client_secret
 API_URL=api_url
 WEB_URL=web_url
+
+GITHUB_TOKEN=github_token # https://github.com/settings/tokens?type=beta

server/data_class.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Literal, Optional
 from pydantic import BaseModel
 
 
@@ -23,3 +23,22 @@ class ExecuteMessage(BaseModel):
 class S3Config(BaseModel):
     s3_bucket: str
     file_path: Optional[str] = None
+
+class GitIssueConfig(BaseModel):
+    repo_name: str
+    page: Optional[int] = None
+    """The page number for paginated results. 
+        Defaults to 1 in the GitHub API."""
+    per_page: Optional[int] = 30
+    """Number of items per page. 
+        Defaults to 30 in the GitHub API."""
+    state: Optional[Literal["open", "closed", "all"]] = 'all'
+    """Filter on issue state. Can be one of: 'open', 'closed', 'all'."""
+    
+    
+class GitDocConfig(BaseModel):
+    repo_name: str
+    file_path: str
+    """File path of the documentation file. eg:'docs/blog/build-ghost.zh-CN.md'"""
+    branch: Optional[str] = 'main'
+    

server/rag/github_file_loader.py

@@ -0,0 +1,70 @@
+
+"""
+This file was originally sourced from the https://github.com/langchain-ai/langchain/blob/master/libs/community/langchain_community/document_loaders/github.py
+and it has been modified based on the requirements provided by petercat.
+"""
+
+import base64
+from abc import ABC
+from typing import Callable, Dict, Iterator, Optional
+import requests
+from langchain_core.documents import Document
+from langchain_core.pydantic_v1 import BaseModel, root_validator
+from langchain_core.utils import get_from_dict_or_env
+
+from langchain_community.document_loaders.base import BaseLoader
+
+
+class BaseGitHubLoader(BaseLoader, BaseModel, ABC):
+    """Load `GitHub` repository Issues. """
+
+    repo: str
+    """Name of repository"""
+    access_token: str
+    """Personal access token - see https://github.com/settings/tokens?type=beta"""
+    github_api_url: str = "https://api.github.com"
+    """URL of GitHub API"""
+
+    @root_validator(pre=True, allow_reuse=True)
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that access token exists in environment."""
+        values["access_token"] = get_from_dict_or_env(
+            values, "access_token", "GITHUB_PERSONAL_ACCESS_TOKEN"
+        )
+        return values
+
+    @property
+    def headers(self) -> Dict[str, str]:
+        return {
+            "Accept": "application/vnd.github+json",
+            "Authorization": f"Bearer {self.access_token}",
+        }
+
+
+class GithubFileLoader(BaseGitHubLoader, ABC):
+    """Load GitHub File"""
+    file_path: str
+    file_extension: str = ".md"
+    branch: str = "main"
+    file_filter: Optional[Callable[[str], bool]]
+
+    def get_file_content_by_path(self, path: str) -> str:
+        base_url = f"{self.github_api_url}/repos/{self.repo}/contents/{path}?ref={self.branch}"
+        response = requests.get(base_url, headers=self.headers)
+        response.raise_for_status()
+
+        if isinstance(response.json(), dict):
+            content_encoded = response.json()["content"]
+            return base64.b64decode(content_encoded).decode("utf-8")
+
+        return ""
+
+    def load(self) -> Iterator[Document]:
+        content = self.get_file_content_by_path(self.file_path)
+
+        metadata = {
+            "path": self.file_path,
+            "source": f"{self.github_api_url}/{self.repo}/blob/"
+            f"{self.branch}/{self.file_path}",
+        }
+        yield Document(page_content=content, metadata=metadata)

server/rag/retrieval.py

@@ -1,18 +1,22 @@
 import json
+from typing import Optional
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import SupabaseVectorStore
 from db.supabase.client import get_client
-from data_class import S3Config
+from data_class import GitDocConfig, GitIssueConfig, S3Config
+from rag.github_file_loader import GithubFileLoader
 from uilts.env import get_env_variable
 
+
 supabase_url = get_env_variable("SUPABASE_URL")
 supabase_key = get_env_variable("SUPABASE_SERVICE_KEY")
+ACCESS_TOKEN=get_env_variable("GITHUB_TOKEN")
 
 
-table_name="antd_knowledge"
-query_name="match_antd_knowledge"
-chunk_size=2000
-
+TABLE_NAME="rag_docs"
+QUERY_NAME="match_rag_docs"
+CHUNK_SIZE=2000
+CHUNK_OVERLAP=20
 
 def convert_document_to_dict(document):
    return document.page_content,
@@ -23,37 +27,84 @@ def init_retriever():
     db = SupabaseVectorStore(
       embedding=embeddings,
       client=get_client(),
-      table_name=table_name,
-      query_name=query_name,
-      chunk_size=chunk_size,
+      table_name=TABLE_NAME,
+      query_name=QUERY_NAME,
+      chunk_size=CHUNK_SIZE,
     )
 
     return db.as_retriever()
 
 
-def add_knowledge(config: S3Config):    
+def init_s3_Loader(config: S3Config):
     from langchain_community.document_loaders import S3DirectoryLoader
+    loader = S3DirectoryLoader(config.s3_bucket, prefix=config.file_path)
+    return loader
+     
+def init_github_issue_loader(config: GitIssueConfig): 
+    from langchain_community.document_loaders import GitHubIssuesLoader
+    
+    loader = GitHubIssuesLoader(
+        repo=config.repo_name,
+        access_token=ACCESS_TOKEN,
+        page=config.page,
+        per_page=config.per_page,
+        state=config.state
+    )
+    return loader
+def init_github_file_loader(config: GitDocConfig): 
+    loader = GithubFileLoader(
+        repo=config.repo_name,
+        access_token=ACCESS_TOKEN,
+        github_api_url="https://api.github.com",
+        branch=config.branch,
+        file_path=config.file_path,
+        file_filter=lambda file_path: file_path.endswith(".md")
+    )
+    return loader
+    
+def supabase_embedding(documents):
     from langchain_text_splitters import CharacterTextSplitter
-
-    try:
-        loader = S3DirectoryLoader(config.s3_bucket, prefix=config.file_path)
-        documents = loader.load()
-        text_splitter = CharacterTextSplitter(chunk_size=2000, chunk_overlap=0)
+    
+    try:    
+        text_splitter = CharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
         docs = text_splitter.split_documents(documents)
         embeddings = OpenAIEmbeddings()
         SupabaseVectorStore.from_documents(
             docs,
             embeddings,
             client=get_client(),
-            table_name=table_name,
-            query_name=query_name,
-            chunk_size=chunk_size,
+            table_name=TABLE_NAME,
+            query_name=QUERY_NAME,
+            chunk_size=CHUNK_SIZE,
         )
         return json.dumps({
             "success": True,
             "message": "Knowledge added successfully!",
             "docs_len": len(documents)
+        }) 
+    except Exception as e:
+        return json.dumps({
+            "success": False,
+            "message": str(e)
         })
+   
+
+def add_knowledge_by_issues(config: GitIssueConfig):    
+    try:
+        loader = init_github_issue_loader(config)
+        documents = loader.load()
+        supabase_embedding(documents)
+    except Exception as e:
+        return json.dumps({
+            "success": False,
+            "message": str(e)
+        })
+   
+def add_knowledge_by_doc(config: GitDocConfig):    
+    try:
+        loader = init_github_file_loader(config)
+        documents = loader.load()
+        supabase_embedding(documents)
     except Exception as e:
         return json.dumps({
             "success": False,

server/routers/rag.py

@@ -1,6 +1,7 @@
+from typing import Optional
 from fastapi import APIRouter, Depends
 from rag import retrieval
-from data_class import S3Config
+from data_class import GitDocConfig, GitIssueConfig, S3Config
 from verify.rate_limit import verify_rate_limit
 
 router = APIRouter(
@@ -10,12 +11,17 @@
 )
 
 
-@router.post("/rag/add_knowledge", dependencies=[Depends(verify_rate_limit)])
-def add_knowledge(config: S3Config):
-    data=retrieval.add_knowledge(config)
+@router.post("/rag/add_knowledge_by_doc")
+def add_knowledge_by_doc(config: GitDocConfig):
+    data=retrieval.add_knowledge_by_doc(config)
     return data
 
-@router.post("/rag/search_knowledge", dependencies=[Depends(verify_rate_limit)])
+@router.post("/rag/add_knowledge_by_issues")
+def add_knowledge_by_issues(config: GitIssueConfig):
+    data=retrieval.add_knowledge_by_issues(config)
+    return data
+
+@router.post("/rag/search_knowledge")
 def search_knowledge(query: str):
     data=retrieval.search_knowledge(query)
     return data

server/sql/rag_docs.sql

@@ -0,0 +1,41 @@
+-- Enable the pgvector extension to work with embedding vectors
+create extension
+if not exists vector;
+
+-- Create a table to store your rag_docs
+create table rag_docs
+(
+  id uuid primary key,
+  content text,
+  -- corresponds to Document.pageContent
+  metadata jsonb,
+  -- corresponds to Document.metadata
+  embedding vector (1536)
+  -- 1536 works for OpenAI embeddings, change if needed
+);
+
+-- Create a function to search for rag_docs
+create function match_rag_docs(
+  query_embedding vector (1536),
+  filter jsonb default '{}'
+) returns table
+(
+  id uuid,
+  content text,
+  metadata jsonb,
+  similarity float
+) language plpgsql as $$
+#variable_conflict use_column
+begin
+  return query
+  select
+    id,
+    content,
+    metadata,
+    1 - (rag_docs.embedding <=> query_embedding
+  ) as similarity
+  from rag_docs
+  where metadata @> filter
+  order by rag_docs.embedding <=> query_embedding;
+end;
+$$;