CodeGPT-Doc.py

import os
import requests
import json
import mimetypes
import logging
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables from a .env file
load_dotenv()

class CodeGPTDocumentManager:
    def __init__(self):
        # Retrieve credentials from environment variables
        self.api_key = os.getenv('CODEGPT_API_KEY')
        self.organization_id = os.getenv('CODEGPT_ORG_ID')
        
        if not self.api_key or not self.organization_id:
            raise ValueError("API key or Organization ID not found in environment variables")

        self.base_url = "https://api.codegpt.co/api/v1"
        self.headers = {
            "accept": "application/json",
            "CodeGPT-Org-Id": self.organization_id,
            "authorization": f"Bearer {self.api_key}",
            "content-type": "application/json"
        }

        # Setup logging
        logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
        self.logger = logging.getLogger(__name__)

    def list_documents(self):
        """List all documents in the storage"""
        try:
            url = f"{self.base_url}/document"
            response = requests.get(url, headers=self.headers)
            if response.status_code == 200:
                data = response.json()
                self.logger.debug(f"List Documents Response: {data}")
                return data
            else:
                self.logger.error(f"Failed to fetch documents. Status: {response.status_code}")
                return None
        except Exception as e:
            self.logger.error(f"Error listing documents: {str(e)}")
            return None

    def delete_document(self, document_id):
        """Delete a document by its ID"""
        try:
            url = f"{self.base_url}/document/{document_id}"
            response = requests.delete(url, headers=self.headers)
            if response.status_code in [200, 204]:
                self.logger.info(f"Deleted document: {document_id}")
                return True
            else:
                self.logger.error(f"Failed to delete document {document_id}. Status: {response.status_code}")
                return False
        except Exception as e:
            self.logger.error(f"Error deleting document {document_id}: {str(e)}")
            return False

    def delete_all_documents(self):
        """Delete all documents"""
        documents = self.list_documents()
        if not documents:
            self.logger.info("No documents to delete.")
            return

        # If the API response wraps the list in a key like 'data', extract it.
        if isinstance(documents, dict) and 'data' in documents:
            docs_list = documents['data']
        else:
            docs_list = documents

        if not docs_list:
            self.logger.info("Document list is empty.")
            return

        for doc in docs_list:
            doc_id = doc.get('id')
            if doc_id:
                success = self.delete_document(doc_id)
                if success:
                    print(f"Deleted document: {doc_id}")
                else:
                    self.logger.error(f"Failed to delete document: {doc_id}")
            else:
                self.logger.warning("Encountered a document without an ID; skipping.")

    def upload_file(self, file_path):
        """Upload a document to CodeGPT API"""
        try:
            url = f"{self.base_url}/document"
            if not os.path.exists(file_path):
                self.logger.error(f"File not found: {file_path}")
                return False

            mime_type = mimetypes.guess_type(file_path)[0] or 'application/octet-stream'
            
            with open(file_path, 'rb') as file:
                files = {'file': (os.path.basename(file_path), file, mime_type)}
                response = requests.post(url, headers=self.headers, files=files)

            if response.status_code in [200, 201]:
                self.logger.info(f"Uploaded file: {file_path}")
                return True
            else:
                self.logger.error(f"Failed to upload file. Status: {response.status_code}, Response: {response.text}")
                return False
        except Exception as e:
            self.logger.error(f"Error uploading file: {str(e)}")
            return False

    def update_metadata(self, document_id, metadata):
        """Update document metadata"""
        try:
            url = f"{self.base_url}/document/{document_id}/metadata"
            response = requests.patch(url, headers=self.headers, data=json.dumps(metadata))
            if response.status_code in [200, 204]:
                self.logger.info(f"Updated metadata for document: {document_id}")
                return True
            else:
                self.logger.error(f"Failed to update metadata for document {document_id}. Status: {response.status_code}")
                return False
        except Exception as e:
            self.logger.error(f"Error updating metadata for document {document_id}: {str(e)}")
            return False

def main():
    """Main execution function"""
    manager = CodeGPTDocumentManager()
    
    while True:
        print("\n=== CodeGPT Document Manager ===")
        print("1. Upload Document")
        print("2. List Documents")
        print("3. Delete Document")
        print("4. Delete All Documents")
        print("5. Update Metadata")
        print("6. Exit")

        choice = input("Enter your choice (1-6): ").strip()
        
        if choice == "1":
            file_path = input("Enter the path to the document: ").strip()
            if manager.upload_file(file_path):
                print("File uploaded successfully")
            else:
                print("Failed to upload file")
        
        elif choice == "2":
            documents = manager.list_documents()
            if documents:
                # Handle the case where documents are wrapped inside a 'data' key.
                if isinstance(documents, dict) and 'data' in documents:
                    docs_list = documents['data']
                else:
                    docs_list = documents
                
                if not docs_list:
                    print("No documents found")
                else:
                    for doc in docs_list:
                        print(f"ID: {doc.get('id')} - Name: {doc.get('name')}")
            else:
                print("No documents found")
        
        elif choice == "3":
            doc_id = input("Enter document ID to delete: ").strip()
            if manager.delete_document(doc_id):
                print("Document deleted successfully")
            else:
                print("Failed to delete document")

        elif choice == "4":
            confirm = input("Are you sure you want to delete ALL documents? (yes/no): ").strip().lower()
            if confirm == "yes":
                manager.delete_all_documents()
                print("Deletion process completed")
            else:
                print("Deletion cancelled")

        elif choice == "5":
            doc_id = input("Enter document ID to update metadata: ").strip()
            title = input("Enter new title: ").strip()
            description = input("Enter description: ").strip()
            summary = input("Enter summary: ").strip()
            keywords = input("Enter keywords (comma separated): ").strip()
            language = input("Enter language code: ").strip()
            
            metadata = {
                "title": title,
                "description": description,
                "summary": summary,
                "keywords": keywords,
                "language": language
            }
            
            if manager.update_metadata(doc_id, metadata):
                print("Metadata updated successfully")
            else:
                print("Failed to update metadata")
        
        elif choice == "6":
            print("Exiting...")
            break
        
        else:
            print("Invalid choice. Please enter a number from 1 to 6.")

if __name__ == "__main__":
    main()