diff --git a/.env.example b/.env.example index 34756db..55371bc 100644 --- a/.env.example +++ b/.env.example @@ -1,5 +1,9 @@ OPENAI_API_KEY="sk-***" LLAMACLOUD_API_KEY="llx-***" + +# Regional Endpoint Configuration (Uncomment the appropriate line for your region) +# LLAMACLOUD_REGION="eu" # Europe + ELEVENLABS_API_KEY="sk_***" pgql_db="postgres" pgql_user="localhost" diff --git a/README.md b/README.md index 8ca3089..c68fe7b 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,11 @@ Next, open the `.env` file and add your API keys: - `ELEVENLABS_API_KEY`: find it [on ElevenLabs Settings](https://elevenlabs.io/app/settings/api-keys) - `LLAMACLOUD_API_KEY`: find it [on LlamaCloud Dashboard](https://cloud.llamaindex.ai?utm_source=demo&utm_medium=notebookLM) +> **🌍 Regional Support**: LlamaCloud operates in multiple regions. If you're using a European region, configure it in your `.env` file: +> +> - For **North America**: This is the default region - no configuration necesary. +> - For **Europe (EU)**: Uncomment and set `LLAMACLOUD_REGION="eu"` + **4. Activate the Virtual Environment** (on mac/unix) diff --git a/src/notebookllama/processing.py b/src/notebookllama/processing.py index a11cf25..c497d99 100644 --- a/src/notebookllama/processing.py +++ b/src/notebookllama/processing.py @@ -1,18 +1,25 @@ +import os +import sys from dotenv import load_dotenv import pandas as pd import json -import os import warnings from datetime import datetime from mrkdwn_analysis import MarkdownAnalyzer from mrkdwn_analysis.markdown_analyzer import InlineParser, MarkdownParser -from llama_cloud_services import LlamaExtract, LlamaParse from llama_cloud_services.extract import SourceText -from llama_cloud.client import AsyncLlamaCloud from typing_extensions import override from typing import List, Tuple, Union, Optional, Dict +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from notebookllama.utils import ( + create_llamacloud_client, + create_llama_extract_client, + create_llama_parse_client, +) + load_dotenv() if ( @@ -20,11 +27,10 @@ and os.getenv("EXTRACT_AGENT_ID", None) and os.getenv("LLAMACLOUD_PIPELINE_ID", None) ): - CLIENT = AsyncLlamaCloud(token=os.getenv("LLAMACLOUD_API_KEY")) - EXTRACT_AGENT = LlamaExtract(api_key=os.getenv("LLAMACLOUD_API_KEY")).get_agent( - id=os.getenv("EXTRACT_AGENT_ID") - ) - PARSER = LlamaParse(api_key=os.getenv("LLAMACLOUD_API_KEY"), result_type="markdown") + CLIENT = create_llamacloud_client() + llama_extract_client = create_llama_extract_client() + EXTRACT_AGENT = llama_extract_client.get_agent(id=os.getenv("EXTRACT_AGENT_ID")) + PARSER = create_llama_parse_client(result_type="markdown") PIPELINE_ID = os.getenv("LLAMACLOUD_PIPELINE_ID") diff --git a/src/notebookllama/querying.py b/src/notebookllama/querying.py index aab1596..cb31f9d 100644 --- a/src/notebookllama/querying.py +++ b/src/notebookllama/querying.py @@ -1,12 +1,16 @@ -from dotenv import load_dotenv import os +import sys +from dotenv import load_dotenv from llama_index.core.query_engine import CitationQueryEngine from llama_index.core.base.response.schema import Response -from llama_index.indices.managed.llama_cloud import LlamaCloudIndex from llama_index.llms.openai import OpenAIResponses from typing import Union, cast +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from notebookllama.utils import create_llamacloud_index + load_dotenv() if ( @@ -16,9 +20,13 @@ ): LLM = OpenAIResponses(model="gpt-4.1", api_key=os.getenv("OPENAI_API_KEY")) PIPELINE_ID = os.getenv("LLAMACLOUD_PIPELINE_ID") - RETR = LlamaCloudIndex( - api_key=os.getenv("LLAMACLOUD_API_KEY"), pipeline_id=PIPELINE_ID - ).as_retriever() + API_KEY = os.getenv("LLAMACLOUD_API_KEY") + + if API_KEY is None or PIPELINE_ID is None: + raise ValueError("LLAMACLOUD_API_KEY and LLAMACLOUD_PIPELINE_ID must be set") + + index = create_llamacloud_index(api_key=API_KEY, pipeline_id=PIPELINE_ID) + RETR = index.as_retriever() QE = CitationQueryEngine( retriever=RETR, llm=LLM, diff --git a/src/notebookllama/server.py b/src/notebookllama/server.py index f798384..b5d8ddc 100644 --- a/src/notebookllama/server.py +++ b/src/notebookllama/server.py @@ -1,9 +1,14 @@ +import os +import sys from querying import query_index from processing import process_file from mindmap import get_mind_map from fastmcp import FastMCP from typing import List, Union, Literal +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + + mcp: FastMCP = FastMCP(name="MCP For NotebookLM") diff --git a/src/notebookllama/utils.py b/src/notebookllama/utils.py new file mode 100644 index 0000000..e07bbd9 --- /dev/null +++ b/src/notebookllama/utils.py @@ -0,0 +1,152 @@ +import os +from typing import Optional, Dict, Any +from llama_cloud.client import AsyncLlamaCloud +from llama_cloud_services import LlamaExtract, LlamaParse +from llama_index.indices.managed.llama_cloud import LlamaCloudIndex + +# LlamaCloud regional endpoints +LLAMACLOUD_REGIONS = { + "default": "https://api.cloud.llamaindex.ai", # North America (default) + "eu": "https://api.cloud.eu.llamaindex.ai", # Europe +} + + +class LlamaCloudConfigError(Exception): + """Raised when LlamaCloud configuration is invalid.""" + + pass + + +def get_llamacloud_base_url() -> Optional[str]: + """ + Get the appropriate LlamaCloud base URL based on region configuration. + + Defaults to North America region and returns the North America endpoint URL + when no region is specified. + + Returns: + str: The base URL for LlamaCloud API + + Raises: + LlamaCloudConfigError: If an invalid region is specified + """ + # Direct base URL override takes precedence + base_url = os.getenv("LLAMACLOUD_BASE_URL") + if base_url: + return base_url + + region = os.getenv("LLAMACLOUD_REGION", "default").lower().strip() + + if region not in LLAMACLOUD_REGIONS: + valid_regions = ", ".join(LLAMACLOUD_REGIONS.keys()) + raise LlamaCloudConfigError( + f"Invalid LLAMACLOUD_REGION '{region}'. Supported regions: {valid_regions}" + ) + + return LLAMACLOUD_REGIONS[region] + + +def get_llamacloud_config() -> Dict[str, Any]: + """ + Get LlamaCloud configuration including base URL. + + Returns: + dict: Configuration dictionary with token and optional base_url + + Raises: + LlamaCloudConfigError: If API key is missing or region is invalid + """ + token = os.getenv("LLAMACLOUD_API_KEY") + if not token: + raise LlamaCloudConfigError( + "LLAMACLOUD_API_KEY environment variable is required" + ) + + config = {"token": token} + + base_url = get_llamacloud_base_url() + if base_url: + config["base_url"] = base_url + + return config + + +def create_llamacloud_client() -> AsyncLlamaCloud: + """ + Create a configured AsyncLlamaCloud client with regional support. + + Returns: + AsyncLlamaCloud: Configured client instance + + Raises: + LlamaCloudConfigError: If API key is missing or region is invalid + """ + config = get_llamacloud_config() + return AsyncLlamaCloud(**config) + + +def create_llama_extract_client() -> LlamaExtract: + """ + Create a configured LlamaExtract client with regional support. + + Returns: + LlamaExtract: Configured client instance + + Raises: + LlamaCloudConfigError: If API key is missing or region is invalid + """ + api_key = os.getenv("LLAMACLOUD_API_KEY") + if not api_key: + raise LlamaCloudConfigError( + "LLAMACLOUD_API_KEY environment variable is required" + ) + + base_url = get_llamacloud_base_url() + return LlamaExtract(api_key=api_key, base_url=base_url) + + +def create_llama_parse_client(result_type: str = "markdown") -> LlamaParse: + """ + Create a configured LlamaParse client with regional support. + + Args: + result_type: The result type for parsing (default: "markdown") + + Returns: + LlamaParse: Configured client instance + + Raises: + LlamaCloudConfigError: If API key is missing or region is invalid + """ + api_key = os.getenv("LLAMACLOUD_API_KEY") + if not api_key: + raise LlamaCloudConfigError( + "LLAMACLOUD_API_KEY environment variable is required" + ) + + base_url = get_llamacloud_base_url() + return LlamaParse(api_key=api_key, result_type=result_type, base_url=base_url) + + +def create_llamacloud_index(api_key: str, pipeline_id: str) -> LlamaCloudIndex: + """ + Create a configured LlamaCloudIndex with regional support. + + Args: + api_key: The API key for authentication + pipeline_id: The pipeline ID to use + + Returns: + LlamaCloudIndex: Configured index instance + + Raises: + LlamaCloudConfigError: If API key or pipeline_id is missing, or region is invalid + """ + if not api_key: + raise LlamaCloudConfigError("API key is required") + + if not pipeline_id: + raise LlamaCloudConfigError("Pipeline ID is required") + + base_url = get_llamacloud_base_url() + return LlamaCloudIndex(api_key=api_key, pipeline_id=pipeline_id, base_url=base_url) diff --git a/tests/test_utils.py b/tests/test_utils.py index 479b862..3c0b384 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,6 +3,7 @@ import pandas as pd from pathlib import Path from dotenv import load_dotenv +from unittest.mock import patch, MagicMock from typing import Callable from pydantic import ValidationError @@ -15,6 +16,16 @@ ) from src.notebookllama.mindmap import get_mind_map from src.notebookllama.models import Notebook +from src.notebookllama.utils import ( + get_llamacloud_base_url, + get_llamacloud_config, + create_llamacloud_client, + create_llama_extract_client, + create_llama_parse_client, + create_llamacloud_index, + LlamaCloudConfigError, + LLAMACLOUD_REGIONS, +) load_dotenv() @@ -189,3 +200,295 @@ def test_images_renaming(images_dir: str): with open(images_dir + "image.png", "wb") as wb: wb.write(bts) os.remove(image) + + +# ============================================================================= +# Regional LlamaCloud Utilities Tests +# ============================================================================= + + +class TestLlamaCloudRegionalUtils: + """Test suite for regional LlamaCloud utility functions.""" + + def test_llamacloud_regions_constant(self): + """Test that LLAMACLOUD_REGIONS contains expected regions.""" + assert "default" in LLAMACLOUD_REGIONS + assert "eu" in LLAMACLOUD_REGIONS + assert LLAMACLOUD_REGIONS["default"] == "https://api.cloud.llamaindex.ai" + assert LLAMACLOUD_REGIONS["eu"] == "https://api.cloud.eu.llamaindex.ai" + + @patch.dict(os.environ, {}, clear=True) + def test_get_llamacloud_base_url_no_region(self): + """Test get_llamacloud_base_url with no region set (defaults to North America).""" + result = get_llamacloud_base_url() + assert result == "https://api.cloud.llamaindex.ai" + + @patch.dict(os.environ, {"LLAMACLOUD_REGION": "eu"}) + def test_get_llamacloud_base_url_eu_region(self): + """Test get_llamacloud_base_url with EU region.""" + result = get_llamacloud_base_url() + assert result == "https://api.cloud.eu.llamaindex.ai" + + @patch.dict(os.environ, {"LLAMACLOUD_REGION": "default"}) + def test_get_llamacloud_base_url_default_region(self): + """Test get_llamacloud_base_url with default region.""" + result = get_llamacloud_base_url() + assert result == "https://api.cloud.llamaindex.ai" + + @patch.dict(os.environ, {"LLAMACLOUD_REGION": "DEFAULT"}) + def test_get_llamacloud_base_url_case_insensitive(self): + """Test get_llamacloud_base_url with case insensitive region.""" + result = get_llamacloud_base_url() + assert result == "https://api.cloud.llamaindex.ai" + + @patch.dict(os.environ, {"LLAMACLOUD_REGION": " default "}) + def test_get_llamacloud_base_url_strips_whitespace(self): + """Test get_llamacloud_base_url strips whitespace from region.""" + result = get_llamacloud_base_url() + assert result == "https://api.cloud.llamaindex.ai" + + @patch.dict(os.environ, {"LLAMACLOUD_BASE_URL": "https://custom.api.com"}) + def test_get_llamacloud_base_url_custom_override(self): + """Test get_llamacloud_base_url with custom base URL override.""" + result = get_llamacloud_base_url() + assert result == "https://custom.api.com" + + @patch.dict( + os.environ, + {"LLAMACLOUD_BASE_URL": "https://custom.api.com", "LLAMACLOUD_REGION": "eu"}, + ) + def test_get_llamacloud_base_url_custom_override_precedence(self): + """Test that custom base URL takes precedence over region.""" + result = get_llamacloud_base_url() + assert result == "https://custom.api.com" + + @patch.dict(os.environ, {"LLAMACLOUD_REGION": "invalid"}) + def test_get_llamacloud_base_url_invalid_region(self): + """Test get_llamacloud_base_url with invalid region raises error.""" + with pytest.raises(LlamaCloudConfigError) as exc_info: + get_llamacloud_base_url() + assert "Invalid LLAMACLOUD_REGION 'invalid'" in str(exc_info.value) + assert "default, eu" in str(exc_info.value) + + @patch.dict(os.environ, {"LLAMACLOUD_API_KEY": "test-key"}, clear=True) + def test_get_llamacloud_config_valid(self): + """Test get_llamacloud_config with valid API key (defaults to North America).""" + result = get_llamacloud_config() + expected = {"token": "test-key", "base_url": "https://api.cloud.llamaindex.ai"} + assert result == expected + + @patch.dict( + os.environ, + {"LLAMACLOUD_API_KEY": "test-key", "LLAMACLOUD_REGION": "eu"}, + clear=True, + ) + def test_get_llamacloud_config_with_region(self): + """Test get_llamacloud_config with region.""" + result = get_llamacloud_config() + expected = { + "token": "test-key", + "base_url": "https://api.cloud.eu.llamaindex.ai", + } + assert result == expected + + @patch.dict(os.environ, {}, clear=True) + def test_get_llamacloud_config_missing_api_key(self): + """Test get_llamacloud_config with missing API key raises error.""" + with pytest.raises(LlamaCloudConfigError): + get_llamacloud_config() + + @patch.dict( + os.environ, + {"LLAMACLOUD_API_KEY": "test-key", "LLAMACLOUD_REGION": "invalid"}, + clear=True, + ) + def test_get_llamacloud_config_invalid_region(self): + """Test get_llamacloud_config with invalid region raises error.""" + with pytest.raises(LlamaCloudConfigError): + get_llamacloud_config() + + @patch("src.notebookllama.utils.AsyncLlamaCloud") + @patch.dict(os.environ, {"LLAMACLOUD_API_KEY": "test-key"}, clear=True) + def test_create_llamacloud_client_valid(self, mock_client_class): + """Test create_llamacloud_client with valid configuration (defaults to North America).""" + mock_instance = MagicMock() + mock_client_class.return_value = mock_instance + + result = create_llamacloud_client() + + mock_client_class.assert_called_once_with( + token="test-key", base_url="https://api.cloud.llamaindex.ai" + ) + assert result == mock_instance + + @patch("src.notebookllama.utils.AsyncLlamaCloud") + @patch.dict( + os.environ, + {"LLAMACLOUD_API_KEY": "test-key", "LLAMACLOUD_REGION": "eu"}, + clear=True, + ) + def test_create_llamacloud_client_with_region(self, mock_client_class): + """Test create_llamacloud_client with region.""" + mock_instance = MagicMock() + mock_client_class.return_value = mock_instance + + result = create_llamacloud_client() + + mock_client_class.assert_called_once_with( + token="test-key", base_url="https://api.cloud.eu.llamaindex.ai" + ) + assert result == mock_instance + + @patch.dict(os.environ, {}, clear=True) + def test_create_llamacloud_client_missing_api_key(self): + """Test create_llamacloud_client with missing API key raises error.""" + with pytest.raises(LlamaCloudConfigError): + create_llamacloud_client() + + @patch("src.notebookllama.utils.LlamaExtract") + @patch.dict(os.environ, {"LLAMACLOUD_API_KEY": "test-key"}, clear=True) + def test_create_llama_extract_client_valid(self, mock_extract_class): + """Test create_llama_extract_client with valid configuration (defaults to North America).""" + mock_instance = MagicMock() + mock_extract_class.return_value = mock_instance + + result = create_llama_extract_client() + + mock_extract_class.assert_called_once_with( + api_key="test-key", base_url="https://api.cloud.llamaindex.ai" + ) + assert result == mock_instance + + @patch("src.notebookllama.utils.LlamaExtract") + @patch.dict( + os.environ, + {"LLAMACLOUD_API_KEY": "test-key", "LLAMACLOUD_REGION": "eu"}, + clear=True, + ) + def test_create_llama_extract_client_with_region(self, mock_extract_class): + """Test create_llama_extract_client with region.""" + mock_instance = MagicMock() + mock_extract_class.return_value = mock_instance + + result = create_llama_extract_client() + + mock_extract_class.assert_called_once_with( + api_key="test-key", base_url="https://api.cloud.eu.llamaindex.ai" + ) + assert result == mock_instance + + @patch.dict(os.environ, {}, clear=True) + def test_create_llama_extract_client_missing_api_key(self): + """Test create_llama_extract_client with missing API key raises error.""" + with pytest.raises(LlamaCloudConfigError): + create_llama_extract_client() + + @patch("src.notebookllama.utils.LlamaParse") + @patch.dict(os.environ, {"LLAMACLOUD_API_KEY": "test-key"}, clear=True) + def test_create_llama_parse_client_default(self, mock_parse_class): + """Test create_llama_parse_client with default parameters (defaults to North America).""" + mock_instance = MagicMock() + mock_parse_class.return_value = mock_instance + + result = create_llama_parse_client() + + mock_parse_class.assert_called_once_with( + api_key="test-key", + result_type="markdown", + base_url="https://api.cloud.llamaindex.ai", + ) + assert result == mock_instance + + @patch("src.notebookllama.utils.LlamaParse") + @patch.dict(os.environ, {"LLAMACLOUD_API_KEY": "test-key"}, clear=True) + def test_create_llama_parse_client_custom_result_type(self, mock_parse_class): + """Test create_llama_parse_client with custom result type (defaults to North America).""" + mock_instance = MagicMock() + mock_parse_class.return_value = mock_instance + + result = create_llama_parse_client(result_type="text") + + mock_parse_class.assert_called_once_with( + api_key="test-key", + result_type="text", + base_url="https://api.cloud.llamaindex.ai", + ) + assert result == mock_instance + + @patch("src.notebookllama.utils.LlamaParse") + @patch.dict( + os.environ, + {"LLAMACLOUD_API_KEY": "test-key", "LLAMACLOUD_REGION": "eu"}, + clear=True, + ) + def test_create_llama_parse_client_with_region(self, mock_parse_class): + """Test create_llama_parse_client with region.""" + mock_instance = MagicMock() + mock_parse_class.return_value = mock_instance + + result = create_llama_parse_client() + + mock_parse_class.assert_called_once_with( + api_key="test-key", + result_type="markdown", + base_url="https://api.cloud.eu.llamaindex.ai", + ) + assert result == mock_instance + + @patch.dict(os.environ, {}, clear=True) + def test_create_llama_parse_client_missing_api_key(self): + """Test create_llama_parse_client with missing API key raises error.""" + with pytest.raises(LlamaCloudConfigError): + create_llama_parse_client() + + @patch("src.notebookllama.utils.LlamaCloudIndex") + @patch.dict(os.environ, {}, clear=True) + def test_create_llamacloud_index_valid(self, mock_index_class): + """Test create_llamacloud_index with valid parameters (defaults to North America).""" + mock_instance = MagicMock() + mock_index_class.return_value = mock_instance + + result = create_llamacloud_index("test-key", "test-pipeline") + + mock_index_class.assert_called_once_with( + api_key="test-key", + pipeline_id="test-pipeline", + base_url="https://api.cloud.llamaindex.ai", + ) + assert result == mock_instance + + @patch("src.notebookllama.utils.LlamaCloudIndex") + @patch.dict(os.environ, {"LLAMACLOUD_REGION": "eu"}, clear=True) + def test_create_llamacloud_index_with_region(self, mock_index_class): + """Test create_llamacloud_index with region.""" + mock_instance = MagicMock() + mock_index_class.return_value = mock_instance + + result = create_llamacloud_index("test-key", "test-pipeline") + + mock_index_class.assert_called_once_with( + api_key="test-key", + pipeline_id="test-pipeline", + base_url="https://api.cloud.eu.llamaindex.ai", + ) + assert result == mock_instance + + @patch.dict(os.environ, {}, clear=True) + def test_create_llamacloud_index_missing_api_key(self): + """Test create_llamacloud_index with missing API key raises error.""" + with pytest.raises(LlamaCloudConfigError) as exc_info: + create_llamacloud_index("", "test-pipeline") + assert "API key is required" in str(exc_info.value) + + @patch.dict(os.environ, {}, clear=True) + def test_create_llamacloud_index_missing_pipeline_id(self): + """Test create_llamacloud_index with missing pipeline ID raises error.""" + with pytest.raises(LlamaCloudConfigError) as exc_info: + create_llamacloud_index("test-key", "") + assert "Pipeline ID is required" in str(exc_info.value) + + @patch.dict(os.environ, {"LLAMACLOUD_REGION": "invalid"}, clear=True) + def test_create_llamacloud_index_invalid_region(self): + """Test create_llamacloud_index with invalid region raises error.""" + with pytest.raises(LlamaCloudConfigError): + create_llamacloud_index("test-key", "test-pipeline") diff --git a/tools/create_llama_cloud_index.py b/tools/create_llama_cloud_index.py index 80a040e..1e3f98f 100644 --- a/tools/create_llama_cloud_index.py +++ b/tools/create_llama_cloud_index.py @@ -1,6 +1,13 @@ +import asyncio import os +import sys + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + + from dotenv import load_dotenv from cli.embedding_app import EmbeddingSetupApp +from src.notebookllama.utils import create_llamacloud_client from llama_cloud import ( PipelineTransformConfig_Advanced, @@ -8,7 +15,6 @@ AdvancedModeTransformConfigSegmentationConfig_Page, PipelineCreate, ) -from llama_cloud.client import LlamaCloud def main(): @@ -16,10 +22,8 @@ def main(): Create a new Llama Cloud index with the given embedding configuration. """ load_dotenv() - client = LlamaCloud(token=os.getenv("LLAMACLOUD_API_KEY")) + client = create_llamacloud_client() - # Run the embedding setup app to get the embedding configuration - # This prompts the user to select an embedding provider and configure the embedding model app = EmbeddingSetupApp() embedding_config = app.run() @@ -45,7 +49,9 @@ def main(): transform_config=transform_config, ) - pipeline = client.pipelines.upsert_pipeline(request=pipeline_request) + pipeline = asyncio.run( + client.pipelines.upsert_pipeline(request=pipeline_request) + ) with open(".env", "a") as f: f.write(f'\nLLAMACLOUD_PIPELINE_ID="{pipeline.id}"') diff --git a/tools/create_llama_extract_agent.py b/tools/create_llama_extract_agent.py index 5908fad..62e688c 100644 --- a/tools/create_llama_extract_agent.py +++ b/tools/create_llama_extract_agent.py @@ -3,15 +3,15 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) -from llama_cloud_services import LlamaExtract from src.notebookllama.models import Notebook +from src.notebookllama.utils import create_llama_extract_client from dotenv import load_dotenv load_dotenv() def main() -> int: - conn = LlamaExtract(api_key=os.getenv("LLAMACLOUD_API_KEY")) + conn = create_llama_extract_client() agent = conn.create_agent(name="q_and_a_agent", data_schema=Notebook) _id = agent.id with open(".env", "a") as f: