efeno/GitPT-Activeloop

1import re2import os3from fastapi import HTTPException4from dotenv import load_dotenv5from llama_index import download_loader6from llama_hub.github_repo import GithubRepositoryReader, GithubClient7from llama_index import VectorStoreIndex8from llama_index.vector_stores import DeepLakeVectorStore9from llama_index.storage.storage_context import StorageContext10import yaml11 12load_dotenv()13 14# Fetch and set API keys15openai_api_key = os.getenv("OPENAI_API_KEY")16 17 18# Check for OpenAI API key19if not openai_api_key:20    raise EnvironmentError("OpenAI API key not found in environment variables")21 22 23def get_validate_token(token_name):24    token = os.getenv(token_name)25    if not token:26        raise EnvironmentError(f"{token_name} not found in environment variables")27    return token28 29 30class InitiazlizeGithubService:31    def __init__(self):32        self.owner = None33        self.repo = None34        self.github_token = get_validate_token("GITHUB_TOKEN")  # Check for GitHub Token35        self.github_client = self.initialize_github_client(self.github_token)36        download_loader("GithubRepositoryReader")37 38    def initialize_github_client(self, github_token):39        return GithubClient(github_token)40 41    def parse_github_url(self, url):42        pattern = r"https://github\.com/([^/]+)/([^/]+)"43        match = re.match(pattern, url)44        return match.groups() if match else (None, None)45 46    def validate_owner_repo(self, owner, repo):47        if bool(owner) and bool(repo):48            self.owner = owner49            self.repo = repo50            return True51 52        return False53 54    def load_repo_data(self, owner, repo, file_type):55        if self.validate_owner_repo(owner, repo):56            loader = GithubRepositoryReader(57                self.github_client,58                owner=self.owner,59                repo=self.repo,60                filter_file_extensions=(61                    [file_type],62                    GithubRepositoryReader.FilterType.INCLUDE,63                ),64                verbose=False,65                concurrent_requests=25,66            )67 68            print(69                f"Loading {self.repo} repository by {self.owner}, file type: {file_type}"70            )71 72            docs = loader.load_data(branch="main")73            print("Documents uploaded:")74            for doc in docs:75                print(doc.metadata)76 77            return docs78 79        else:80            raise HTTPException(81                status_code=400,82                detail="Invalid GitHub URL. Please enter a valid GitHub URL",83            )84 85 86class InitiazlizeActiveloopService:87    def __init__(self):88        self.active_loop_token = get_validate_token(89            "ACTIVELOOP_TOKEN"90        )  # Check for Activeloop Token91        self.dataset_path = self.get_user_info("dataset_path")92        self.vector_store = DeepLakeVectorStore(93            dataset_path=f"hub://{self.dataset_path}",94            overwrite=True,95            runtime={"tensor_db": True},96        )97 98        self.storage_context = StorageContext.from_defaults(99            vector_store=self.vector_store100        )101 102    def upload_to_activeloop(self, docs):103        self.index = VectorStoreIndex.from_documents(104            docs, storage_context=self.storage_context105        )106        self.query_engine = self.index.as_query_engine()107 108    def get_user_info(self, user_info):109        with open("resources.yaml", "r") as file:110            yaml_data = yaml.safe_load(file)111 112        retrieved_info = yaml_data["info"][user_info]113        return retrieved_info114