Allow uploading attachments

api/src/env_var.py

@@ -15,3 +15,5 @@
 
 CSE_ID = environ["CSE_ID"]
 CSE_API_KEY = environ["GOOGLE_API_KEY"]
+
+PROD_ENV = environ.get("ENV", "dev") == "prod"

api/src/main.py

@@ -1,4 +1,4 @@
-import asyncio
+from io import BytesIO
 from time import time
 from typing import Any, Callable, Generator
 
@@ -7,7 +7,7 @@
 from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi_utilities import add_timer_middleware
 
-from .services.storage import StorageManager
+from .services.storage import BLOB_BASE_URI, StorageManager, UploadItemParams
 from .utils.chat_request import chat_request
 from .utils.chat_utils import (
     ContentCategory,
@@ -24,14 +24,18 @@
     GetCustomSourcesRequest,
     generate_custom_source,
 )
+from .utils.custom_sources.manage_attachments import ManageAttachments
+from .utils.custom_sources.read_content import ReadContent
 from .utils.custom_sources.save_copied_source import CopiedPasteSourceRequest, save_copied_source
 from .utils.custom_sources.save_uploaded_sources import UploadedFiles
+from .utils.decorators.retry_decorator import RetryConfig, retry
 from .utils.detect_content_category import DetectContentCategoryRequest, detect_content_category
 from .utils.generate_audiocast import GenerateAudioCastRequest, GenerateAudiocastException, generate_audiocast
 from .utils.generate_audiocast_source import GenerateAudiocastSource, generate_audiocast_source
 from .utils.get_audiocast import get_audiocast
 from .utils.get_session_title import GetSessionTitleModel, get_session_title
 from .utils.session_manager import SessionManager, SessionModel
+from .utils.summarize_custom_sources import SummarizeCustomSourcesRequest, summarize_custom_sources
 
 app = FastAPI(title="Audiora", version="1.0.0")
 
@@ -69,26 +73,35 @@ def root():
 
 
 @app.post("/chat/{session_id}", response_model=Generator[str, Any, None])
-def chat_endpoint(
+async def chat_endpoint(
     session_id: str,
     request: SessionChatRequest,
     background_tasks: BackgroundTasks,
 ):
     """Chat endpoint"""
     category = request.contentCategory
+    attachments = request.attachments
+
     db = SessionManager(session_id, category)
     db._add_chat(request.chatItem)
 
+    attachment_manager = ManageAttachments(session_id)
+    sources_summary = await attachment_manager.get_attachments_summary(db, attachments)
+
     def on_finish(text: str):
         background_tasks.add_task(db._update, {"status": "collating"})
         background_tasks.add_task(
             db._add_chat,
             SessionChatItem(role="assistant", content=text),
         )
 
+        if attachments:
+            background_tasks.add_task(attachment_manager.store_attachments, attachments)
+
     response = chat_request(
         content_category=category,
         previous_messages=db._get_chats(),
+        reference_material=sources_summary,
         on_finish=on_finish,
     )
 
@@ -128,27 +141,22 @@ async def get_signed_url_endpoint(blobname: str):
     """
     Get signed URL for generated audiocast
     """
-    retry_count = 0
-    max_retries = 3
-    errors: list[str] = []
 
-    while retry_count < max_retries:
-        try:
-            url = StorageManager().get_signed_url(blobname=blobname)
-            return JSONResponse(
-                content=url,
-                headers={
-                    "Content-Type": "application/json",
-                    "Cache-Control": "public, max-age=86390, immutable",
-                },
-            )
-        except Exception as e:
-            errors.append(str(e))
+    @retry(RetryConfig(max_retries=3, delay=5, backoff=1.5))
+    def handler() -> str | None:
+        return StorageManager().get_signed_url(blobname=blobname)
 
-        await asyncio.sleep(5)
-        retry_count += 1
+    url = handler()
+    if not url:
+        raise HTTPException(status_code=500, detail="Failed to get signed URL")
 
-    raise HTTPException(status_code=500, detail="".join(errors))
+    return JSONResponse(
+        content=url,
+        headers={
+            "Content-Type": "application/json",
+            "Cache-Control": "public, max-age=86390, immutable",
+        },
+    )
 
 
 @app.post("/get-session-title", response_model=str)
@@ -216,3 +224,42 @@ async def detect_category_endpoint(request: DetectContentCategoryRequest):
     Detect category of a given content
     """
     return await detect_content_category(request.content)
+
+
+@app.post("/store-file-upload", response_model=str)
+async def store_file_upload(file: UploadFile, filename: str = Form(...), preserve: bool = Form(False)):
+    """
+    Store file uploaded from the frontend
+    """
+    print(f"Storing file: {filename}. Preserve: {preserve}")
+
+    storage_manager = StorageManager()
+    file_exists = storage_manager.check_blob_exists(filename)
+    if file_exists:
+        return storage_manager.get_gcs_url(filename)
+
+    file_content = await ReadContent()._read_file(file, preserve=preserve)
+    content_type = (
+        file.content_type or "application/octet-stream"
+        if preserve or isinstance(file_content, BytesIO)
+        else "text/plain"
+    )
+
+    result = storage_manager.upload_to_gcs(
+        item=file_content,
+        blobname=f"{BLOB_BASE_URI}/{filename}",
+        params=UploadItemParams(
+            cache_control="public, max-age=31536000",
+            content_type=content_type,
+        ),
+    )
+
+    return result
+
+
+@app.post("/summarize-custom-sources", response_model=str)
+async def summarize_custom_sources_endpoint(request: SummarizeCustomSourcesRequest):
+    """
+    Summarize custom sources from specified source URLs
+    """
+    return await summarize_custom_sources(request.sourceURLs)

api/src/services/storage.py

@@ -31,6 +31,8 @@ class UploadItemParams:
 
 
 class StorageManager:
+    bucket_name = BUCKET_NAME
+
     def check_blob_exists(self, filename: str, root_path=BLOB_BASE_URI):
         """check if a file exists in the bucket"""
         blobname = f"{root_path}/{filename}"
@@ -123,3 +125,16 @@ def get_signed_url(self, blobname, expiration=datetime.timedelta(days=1)):
                 expiration=expiration,
                 method="GET",
             )
+
+    def get_gcs_url(self, filename: str):
+        """get full path to a file in the bucket"""
+        blobname = f"{BLOB_BASE_URI}/{filename}"
+        return f"gs://{BUCKET_NAME}/{blobname}"
+
+    def get_blob(self, blobname: str):
+        """get a blob object"""
+        return bucket.blob(blobname)
+
+    def get_blobname_from_url(self, url: str):
+        """get blobname from a URL"""
+        return url.replace(f"gs://{self.bucket_name}/", "")

api/src/utils/chat_request.py

@@ -4,19 +4,22 @@
 from src.utils.chat_utils import ContentCategory, SessionChatItem
 
 
-def get_system_message(content_category: ContentCategory):
+def get_system_message(content_category: ContentCategory, reference_material: str | None = None):
     return f"""
     1. You're a super-intelligent AI. Your task is to understand what audiocast a user wants to listen to.
     2. You will steer the conversation providing eliciting questions until you have enough context.
-    3. Keep the conversation exchange short, say 3-5 back and forth i.e., questions and answers.
+    3. If the user provides a reference material, steer the conversation based on it until you have enough context to understand what audiocast the user wants.
+    4. Keep the conversation exchange short, say 3-5 back and forth i.e., questions and answers.
     4. As soon as you have enough context and the user's request is clear terminate the conversation by saying "Ok, thanks for clarifying! You want to listen to [Best case summary of user request so far]. Please click the button below to start generating the audiocast."
     6. If the user's request remains unclear after 5 responses for clarity, terminate the conversation by saying "Your request is not very specific but from what I understand, you want to listen to [Best case summary of user request so far]. Please click the button below to start generating the audiocast."
 
+    {"REFERENCE MATERIAL: " + reference_material if reference_material else ""}
 
     GENERAL IDEA AND WORKFLOW:
     1. A user comes to you with a request for an audiocast of type {content_category}.
-    2. You need to ask the user questions (elicitation) to understand what kind of audiocast they want to listen to.
-    3. Once you have enough context, within 3-5 exchanges, you should terminate the conversation.
+    2. The request can include a reference material: a high-level description of the audiocast they want.
+    3. You will ask the user questions (elicitation) to understand what kind of audiocast they want to listen to.
+    4. Once you have enough context, within 3-5 exchanges, you should terminate the conversation.
 
     IMPORTANT NOTES:
     1. Your task is to understand the user's request only by eliciting questions.
@@ -28,12 +31,13 @@ def get_system_message(content_category: ContentCategory):
 def chat_request(
     content_category: ContentCategory,
     previous_messages: List[SessionChatItem],
+    reference_material: Optional[str] = None,
     on_finish: Optional[Callable[[str], Any]] = None,
 ):
     response_stream = get_openai().chat.completions.create(
         model="gpt-4o",
         messages=[
-            {"role": "system", "content": get_system_message(content_category)},
+            {"role": "system", "content": get_system_message(content_category, reference_material)},
             *[
                 {"role": "user", "content": msg.content}
                 if msg.role == "user"

api/src/utils/chat_utils.py

@@ -1,5 +1,5 @@
 import uuid
-from typing import Dict, List, Literal
+from typing import Dict, List, Literal, Optional
 
 from pydantic import BaseModel, Field
 
@@ -57,3 +57,4 @@ class SessionChatItem(BaseModel):
 class SessionChatRequest(BaseModel):
     contentCategory: ContentCategory
     chatItem: SessionChatItem
+    attachments: Optional[List[str]] = None

api/src/utils/custom_sources/base_utils.py

@@ -1,6 +1,7 @@
 from typing import Literal, Optional, TypedDict, cast
 
 from google.cloud.firestore_v1 import DocumentReference
+from google.cloud.firestore_v1.base_query import FieldFilter
 from pydantic import BaseModel
 
 from src.services.firestore_sdk import (
@@ -103,3 +104,19 @@ def _get_custom_sources(self) -> list[CustomSourceModelDict]:
 
     def _delete_custom_source(self, source_id: str):
         return self._get_doc_ref(source_id).delete()
+
+    def _get_custom_source_by_url(self, url: str):
+        self._check_document()
+        try:
+            session_ref = self._get_collection(self.collection).document(self.doc_id)
+
+            query = session_ref.collection(self.sub_collection).where(filter=FieldFilter("url", "==", url))
+
+            docs = query.get()
+
+            for doc in docs:
+                if doc.exists:
+                    return cast(CustomSourceModel, self._safe_to_dict(doc.to_dict()))
+        except Exception as e:
+            print(f"Error getting custom sources for Session: {self.doc_id}", e)
+        return None

api/src/utils/custom_sources/extract_url_content.py

@@ -5,7 +5,8 @@
 from bs4 import BeautifulSoup, Tag
 from pydantic import BaseModel
 
-from src.utils.decorators import process_time
+from src.services.storage import StorageManager
+from src.utils.decorators.base import process_time
 
 from .base_utils import SourceContent
 from .read_content import ReadContent
@@ -42,14 +43,24 @@ def _extract_html(self, content: bytes) -> tuple[str, dict]:
 
         return self._clean_text(text_content), metadata
 
+    def _resolve_gcs_url(self, url) -> str:
+        if url.startswith("gs://"):
+            storage_manager = StorageManager()
+            blobame = storage_manager.get_blobname_from_url(url)
+            return storage_manager.get_signed_url(blobame)
+
+        return url
+
     @process_time()
     def _extract(self, url: str) -> SourceContent:
-        parsed_url = urlparse(url)
+        resolved_url = self._resolve_gcs_url(url)
+
+        parsed_url = urlparse(resolved_url)
         if not parsed_url.scheme or not parsed_url.netloc:
             raise ValueError("Invalid URL provided")
 
         try:
-            response = httpx.get(url)
+            response = httpx.get(resolved_url)
             response.raise_for_status()
             content_type = response.headers.get("content-type", "").lower()
 

api/src/utils/custom_sources/generate_url_source.py

@@ -20,7 +20,10 @@ class DeleteCustomSourcesRequest(BaseModel):
     sourceId: str
 
 
-def generate_custom_source(request: GenerateCustomSourceRequest, background_tasks: BackgroundTasks):
+def generate_custom_source(
+    request: GenerateCustomSourceRequest,
+    background_tasks: BackgroundTasks | None = None,
+):
     extractor = ExtractURLContent()
     content = extractor._extract(request.url)
 
@@ -33,5 +36,9 @@ def save_to_firestore():
         manager = CustomSourceManager(request.sessionId)
         manager._set_custom_source(custom_source)
 
-    background_tasks.add_task(save_to_firestore)
+    if background_tasks:
+        background_tasks.add_task(save_to_firestore)
+    else:
+        save_to_firestore()
+
     return content.model_dump()

api/src/utils/custom_sources/manage_attachments.py

@@ -0,0 +1,47 @@
+import asyncio
+
+from src.utils.decorators.base import use_cache_manager
+from src.utils.make_seed import get_hash
+from src.utils.session_manager import SessionManager
+from src.utils.summarize_custom_sources import summarize_custom_sources
+
+from .base_utils import CustomSourceManager
+from .generate_url_source import GenerateCustomSourceRequest, generate_custom_source
+
+
+class ManageAttachments:
+    def __init__(self, session_id: str):
+        self.session_id = session_id
+
+    async def get_attachments_summary(self, db: SessionManager, attachments: list[str] | None):
+        """
+        Manage custom sources uploaded by the user
+        """
+        sources_summary: str | None = None
+        if attachments:
+            attachments.sort(key=lambda x: x.lower())
+
+            @use_cache_manager(get_hash(attachments))
+            async def handler():
+                summary = await summarize_custom_sources(attachments)
+                db._update_source(summary)
+                return summary
+
+            sources_summary = await handler()
+
+        return sources_summary
+
+    async def store_attachments(self, attachments: list[str]):
+        """
+        Store attachments as custom sources of type links
+        """
+        cs_manager = CustomSourceManager(self.session_id)
+
+        async def _handler(url: str):
+            custom_source = cs_manager._get_custom_source_by_url(url)
+            if not custom_source:
+                request = GenerateCustomSourceRequest(url=url, sessionId=self.session_id)
+                return generate_custom_source(request)
+
+        await asyncio.gather(*[_handler(url) for url in attachments], return_exceptions=True)
+        return True