Adding Compass Tool (#240)

* Initial commit for Compass Tool * coral-web: update first turn suggestions (#225) * Addressed comments on the previous commit * Added custom_context to parse and filters to search * Formatting changes * feat(assistant): Update conversation header (#229) * feat(assistant): Update conversation header * feat(assistant): Update conversation header * update tooltip * fix IconButton import * Setup: fix dependencies (#243) * feat(toolkit): show conversation list on with agents list (#231) * Setup: remove redundant dependency (#244) * Improve Auth guide (#235) improve auth guide * coral-web: create agent + agent base form (#212) * create new agent form * render agents list * link back to base agent page * perform post submission actions * move input labels outside of input borders * move submit button * gen client * factor out agent base form * use default content type * list custom agents * check if agent name is unique * and deployment + env variables to agent form * add deployments * clean up * add submit modal * push to new agent page on create success * only show available deployments * remove deployments dropdown * feat: List Assistants returned from API (#242) * feat: list agents from BE * Update src/interfaces/coral_web/src/components/Conversation/MessagingContainer.tsx Co-authored-by: misspia-cohere <140425731+misspia-cohere@users.noreply.github.com> --------- Co-authored-by: misspia-cohere <140425731+misspia-cohere@users.noreply.github.com> * nits --------- Co-authored-by: Khalil Najjar <knajjars@gmail.com> * Propagate should_store logic to the File upload logic (#247) should_store logic to the attach_files_to_messages * [backend] passing in agent when calling streaming chat (#237) * initial chat changes * need to test * done * Remove old test * default model for agents right now * Organizations DB models, CRUD and tests (#238) * Organizations initial commit * Organizations initial commit * Organizations initial commit - sync main * Organizations initial commit - sync main - blacked * Organizations initial commit - sync main - lint * Organizations initial commit - sync main - tests * Organizations initial commit - sync main - tests * Organizations initial commit - review fixes * Organizations initial commit - review fixes * Split up display name and name for Tools (#241) * CHange * change * coral-web: update assistant (#248) * create new agent form * render agents list * link back to base agent page * perform post submission actions * move input labels outside of input borders * move submit button * gen client * factor out agent base form * check if agent name is unique * and deployment + env variables to agent form * add deployments * clean up * remove deployments dropdown * init agent page * create agent drawer * use agent form * add update agent request * add get agent request * fix rebase errors * remove model from agent form * add agent drawer to conversation * remove unused componenets * use new agent name for update success message * fix(toolkit): Address code feedback + improvements (#249) address code feedback improvements --------- Co-authored-by: Khalil Najjar <knajjars@gmail.com> * feat(toolkit): add/remove recently used agents (#250) * feat(toolkit): add/remove recently used agents * merge main * Add a walkthrough guide of the toolkit (#251) * GUide * Chang * Change * Change * Update docs/walkthrough/walkthrough.md Co-authored-by: Luísa Moura <luisa@cohere.com> * Update walkthrough.md --------- Co-authored-by: Luísa Moura <luisa@cohere.com> * coral-web: fix agent info panel opening by default (#253) cast isEditAgentPanelOpen to boolean * [backend] enforce agent update with user-id (#246) * updates * remove client changes * remove logs * use better header user id check * fix validators * typo * Metrics: add middleware (#185) * Metrics: add middleware * add chat calls * merge * lint * make it async * add user id * add more fields * add retry and duration * add meta * comments * fix tests * improve error handling * rename fields * match spec * comments * clean code * only create loop when theres endpoint * add assistant id to chat * feat(toolkit): show assistant welcome message (#255) * feat(toolkit): show assistant welcome message * feat(toolkit): show assistant welcome message * frontend: Login, logout, and account creation (#179) * add login page components * Add Register page and hooks for auth * Add the register page and connect all the frontend elements * Redirect to /login if the token expires and clean up some console errors * Add error messages for failed logins * frontend: Add Single Sign-on to Toolkit (#227) * Add Google SSO login plus OpenID components * Dynamically set SSO login buttons and show or hide username and password based on auth_strategies --------- Co-authored-by: Tianjing Li <tianjinglimail@gmail.com> * fix startup event * Fix build errors and update the API client * Fix tests by adding missing env vars Add test OIDC_WELL_KNOWN_ENDPOINT var to fixtures * Add a walkthrough guide of the toolkit (#251) * GUide * Chang * Change * Change * Update docs/walkthrough/walkthrough.md Co-authored-by: Luísa Moura <luisa@cohere.com> * Update walkthrough.md --------- Co-authored-by: Luísa Moura <luisa@cohere.com> * coral-web: fix agent info panel opening by default (#253) cast isEditAgentPanelOpen to boolean * [backend] enforce agent update with user-id (#246) * updates * remove client changes * remove logs * use better header user id check * fix validators * typo * Metrics: add middleware (#185) * Metrics: add middleware * add chat calls * merge * lint * make it async * add user id * add more fields * add retry and duration * add meta * comments * fix tests * improve error handling * rename fields * match spec * comments * clean code * only create loop when theres endpoint * add assistant id to chat * feat(toolkit): show assistant welcome message (#255) * feat(toolkit): show assistant welcome message * feat(toolkit): show assistant welcome message --------- Co-authored-by: misspia-cohere <pia@cohere.com> Co-authored-by: Tianjing Li <tianjinglimail@gmail.com> Co-authored-by: Beatrix De Wilde <128378696+BeatrixCohere@users.noreply.github.com> Co-authored-by: Luísa Moura <luisa@cohere.com> Co-authored-by: misspia-cohere <140425731+misspia-cohere@users.noreply.github.com> Co-authored-by: Scott <146760070+scott-cohere@users.noreply.github.com> Co-authored-by: Khalil Najjar <knajjars@gmail.com> * fix(toolkit): UX/UI improvements (#257) * Add error troubleshooting (#262) * Update chat.py --------- Co-authored-by: Jessica Wu <jessica@cohere.ai> Co-authored-by: Khalil Najjar <knajjars@gmail.com> Co-authored-by: Luísa Moura <luisa@cohere.com> Co-authored-by: Tomeu <tomeu@cohere.com> Co-authored-by: Tianjing Li <tianjinglimail@gmail.com> Co-authored-by: misspia-cohere <140425731+misspia-cohere@users.noreply.github.com> Co-authored-by: Eugene P <144219719+EugeneLightsOn@users.noreply.github.com> Co-authored-by: Scott <146760070+scott-cohere@users.noreply.github.com> Co-authored-by: Beatrix De Wilde <128378696+BeatrixCohere@users.noreply.github.com> Co-authored-by: Alex W <alex@blinkenlights.ca> Co-authored-by: misspia-cohere <pia@cohere.com>
cohere-ai · Jun 21, 2024 · 0c5db02 · 0c5db02
1 parent eb34bf8
commit 0c5db02
Showing 1 changed file with 202 additions and 0 deletions.
diff --git a/src/backend/tools/compass.py b/src/backend/tools/compass.py
@@ -0,0 +1,202 @@
+import logging
+import os
+from typing import Any, Dict, List
+
+from compass_sdk import MetadataConfig, ParserConfig
+from compass_sdk.compass import CompassClient
+from compass_sdk.parser import CompassParserClient
+
+from backend.tools.base import BaseTool
+
+logger = logging.getLogger()
+
+
+class CompassTool(BaseTool):
+    """Tool to interact with a Compass instance."""
+
+    @classmethod
+    def is_available(cls) -> bool:
+        vars = [
+            "COHERE_COMPASS_URL",
+            "COHERE_COMPASS_USERNAME",
+            "COHERE_COMPASS_PASSWORD",
+        ]
+        return all(os.getenv(var) is not None for var in vars)
+
+    def __init__(
+        self,
+        compass_url=None,
+        compass_username=None,
+        compass_password=None,
+        metadata_config=MetadataConfig(),
+        parser_config=ParserConfig(),
+    ):
+        """Initialize the Compass tool. Pass the Compass URL, username, and password
+        as arguments or as environment variables."""
+        self.url = compass_url or os.getenv("COHERE_COMPASS_URL")
+        self.username = compass_username or os.getenv("COHERE_COMPASS_USERNAME")
+        self.password = compass_password or os.getenv("COHERE_COMPASS_PASSWORD")
+        self.parser_config = parser_config
+        self.metadata_config = metadata_config
+        try:
+            # Try initializing Compass Parser and Client and call list_indexes
+            # to check if the credentials are correct.
+            self.parser_client = CompassParserClient(
+                parser_url=self.url + "/parse",
+                username=self.username,
+                password=self.password,
+                parser_config=self.parser_config,
+                metadata_config=self.metadata_config,
+            )
+            self.compass_client = CompassClient(
+                index_url=self.url,
+                username=self.username,
+                password=self.password,
+            )
+            self.compass_client.list_indexes()
+        except Exception as e:
+            logger.exception(f"Compass Tool: Error initializing Compass client: {e}")
+            raise e
+
+    def call(self, parameters: dict, **kwargs: Any) -> List[Dict[str, Any]]:
+        """Call the Compass tool. Allowed `action` values:
+        - list_indexes: List all indexes in Compass.
+        - create_index: Create a new index in Compass.
+        - delete_index: Delete an existing index in Compass.
+        - create: Create a new document in Compass.
+        - search: Search for documents in Compass.
+        - update: Update an existing document in Compass.
+        - delete: Delete an existing document in Compass.
+        """
+        # Check if action is specified
+        if not parameters.get("action", None):
+            logger.error(
+                "Compass Tool: No action specified. "
+                "No action will be taken. "
+                f"Parameters specified: {parameters}"
+            )
+            return
+        # Check if action is valid
+        valid_actions = [
+            "list_indexes",
+            "create_index",
+            "delete_index",
+            "create",
+            "search",
+            "update",
+            "delete",
+        ]
+        if parameters["action"] not in valid_actions:
+            logger.error(
+                f"Compass Tool: Invalid action {parameters['action']}. "
+                "No action will be taken. "
+                f"Parameters specified: {parameters}"
+            )
+            return
+        # Check if index is specified
+        if not parameters.get("index", None) and parameters["action"] != "list_indexes":
+            logger.error(
+                "Compass Tool: No index specified. "
+                "No action will be taken. "
+                f"Parameters specified: {parameters}"
+            )
+            return
+        # Index-related actions
+        if parameters["action"] == "list_indexes":
+            return self.compass_client.list_indexes()
+        elif parameters["action"] == "create_index":
+            return self.compass_client.create_index(index_name=parameters["index"])
+        elif parameters["action"] == "delete_index":
+            return self.compass_client.delete_index(index_name=parameters["index"])
+        # Check if file_id is specified for file-related actions
+        if not parameters.get("file_id", None):
+            logger.error(
+                "Compass Tool: No uninque identifier file_id specified. "
+                "No action will be taken. "
+                f"Parameters specified: {parameters}"
+            )
+            return
+        # Create index if it does not exist
+        self.compass_client.create_index(index_name=parameters.get("index"))
+        # Perform the file-related action
+        if parameters["action"] == "create":
+            self._create(parameters, **kwargs)
+        elif parameters["action"] == "search":
+            self._search(parameters, **kwargs)
+        elif parameters["action"] == "update":
+            self._update(parameters, **kwargs)
+        elif parameters["action"] == "delete":
+            self._delete(parameters, **kwargs)
+
+    def _parse(self, parameters: dict, **kwargs: Any) -> None:
+        """Parse the input file."""
+        if "file_path" not in parameters:
+            logger.error(
+                "Compass Tool: No file_path specified for "
+                "create/update operation. "
+                "No action will be taken. "
+                f"Parameters specified: {parameters}"
+            )
+            return None
+        file_path = parameters["file_path"]
+        if not os.path.exists(file_path):
+            logger.error(
+                f"Compass Tool: File {file_path} does not exist. "
+                "No action will be taken."
+                f"Parameters specified: {parameters}"
+            )
+            return None
+        parser_config = self.parser_config or parameters.get("parser_config", None)
+        metadata_config = self.metadata_config or parameters.get(
+            "metadata_config", None
+        )
+        return self.parser_client.process_file(
+            file_path=file_path,
+            file_id=parameters["file_id"],
+            parser_config=parser_config,
+            metadata_config=metadata_config,
+            custom_context=parameters.get("custom_context", None),
+            is_dataset=False,
+        )
+
+    def _create(self, parameters: dict, **kwargs: Any) -> Dict[str, str]:
+        """Insert the document into Compass"""
+        compass_docs = self._parse(parameters, **kwargs)
+        if compass_docs is None:
+            # Parsing failed
+            return
+        error = self.compass_client.insert_docs(
+            index_name=parameters["index"], docs=compass_docs
+        )
+        if error is not None:
+            logger.error(
+                f"Compass Tool: Error inserting/updating document "
+                f"into Compass: {error}"
+            )
+
+    def _search(self, parameters: dict, **kwargs: Any) -> None:
+        """Run a search query on Compass and return the
+        top_k results. By default, k=10."""
+        if not parameters.get("query", None):
+            logger.error(
+                "Compass Tool: No search query specified. ",
+                "Returning empty list. " "Parameters specified: {parameters}",
+            )
+            return []
+        return self.compass_client.search(
+            index_name=parameters["index"],
+            query=parameters["query"],
+            top_k=parameters.get("top_k", 10),
+            filters=parameters.get("filters", None),
+        )
+
+    def _update(self, parameters: dict, **kwargs: Any) -> None:
+        """Update file in Compass"""
+        self._delete(parameters, **kwargs)
+        self._create(parameters, **kwargs)
+
+    def _delete(self, parameters: dict, **kwargs: Any) -> None:
+        """Delete file from Compass"""
+        self.compass_client.delete_document(
+            index_name=parameters["index"], doc_id=parameters["file_id"]
+        )