Skip to content

Commit

Permalink
saving
Browse files Browse the repository at this point in the history
  • Loading branch information
ClaytonSmith committed Jun 14, 2024
1 parent 70dd6ec commit 841355d
Show file tree
Hide file tree
Showing 4 changed files with 212 additions and 26 deletions.
68 changes: 45 additions & 23 deletions docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@ services:
- '5432:5432'
volumes:
- db:/var/lib/postgresql/data
networks:
- my_network

test_db:
image: postgres:14.1-alpine
restart: unless-stopped
Expand All @@ -26,40 +29,50 @@ services:
- POSTGRES_DB=test_toolkit
ports:
- '5433:5432'
networks:
- my_network

backend:
build:
context: .
args:
INSTALL_COMMUNITY_DEPS: false
dockerfile: ./src/backend/Dockerfile
# develop:
# watch:
# - action: sync
# path: ./src/backend
# target: /workspace/src/backend
# ignore:
# - __pycache__/
# - alembic/
# - data/
# - action: sync
# path: ./src/community
# target: /workspace/src/community
# ignore:
# - __pycache__/
# - alembic/
# - data/
develop:
watch:
- action: sync
path: ./src/backend
target: /workspace/src/backend
ignore:
- __pycache__/
- alembic/
- data/
- action: sync
path: ./src/community
target: /workspace/src/community
ignore:
- __pycache__/
- alembic/
- data/
stdin_open: true
tty: true
ports:
- '8000:8000'
# - '8001:8001'
depends_on:
- db
# extra_hosts:
# - "host.docker.internal:host-gateway" # Map host.docker.internal to the host's IP
environment:
- HOST_SERVICE_URL=http://host.docker.internal:8081 # Environment variable for host service URL
volumes:
# Mount alembic folder to sync migrations
- ./src/backend/alembic:/workspace/src/backend/alembic
# Mount data folder to sync uploaded files
- ./src/backend/data:/workspace/src/backend/data
# network_mode: host
networks:
- my_network

frontend:
build:
Expand All @@ -72,20 +85,29 @@ services:
restart: always
ports:
- 4000:4000
# develop:
# watch:
# - action: sync
# path: ./src/interfaces/coral_web
# target: /app
# ignore:
# - node_modules/
develop:
watch:
- action: sync
path: ./src/interfaces/coral_web
target: /app
ignore:
- node_modules/
networks:
- my_network
terrarium:
image: ghcr.io/cohere-ai/terrarium:latest
ports:
- '8080:8080'
expose:
- '8080'
networks:
- my_network

volumes:
db:
name: cohere_toolkit_db
driver: local

networks:
my_network:
driver: bridge
25 changes: 22 additions & 3 deletions src/backend/config/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,14 @@
ReadFileTool,
SearchFileTool,
TavilyInternetSearch,
LangChainMinimapRetriever
)

"""
List of available tools. Each tool should have a name, implementation, is_visible and category.
List of available tools. Each tool should have a name, implementation, is_visible and category.
They can also have kwargs if necessary.
You can switch the visibility of a tool by changing the is_visible parameter to True or False.
You can switch the visibility of a tool by changing the is_visible parameter to True or False.
If a tool is not visible, it will not be shown in the frontend.
If you want to add a new tool, check the instructions on how to implement a retriever in the documentation.
Expand All @@ -32,6 +33,7 @@ class ToolName(StrEnum):
Python_Interpreter = "Python_Interpreter"
Calculator = "Calculator"
Tavily_Internet_Search = "Internet_Search"
Minimap = "Minimap"


ALL_TOOLS = {
Expand All @@ -52,6 +54,23 @@ class ToolName(StrEnum):
category=Category.DataLoader,
description="Retrieves documents from Wikipedia using LangChain.",
),
ToolName.Minimap: ManagedTool(
name=ToolName.Minimap,
implementation=LangChainMinimapRetriever,
parameter_definitions={
"query": {
"description": "Query for retrieval.",
"type": "str",
"required": True,
}
},
is_visible=True,
is_available=LangChainMinimapRetriever.is_available(),
error_message="LangChainMinimapRetriever not available.",
category=Category.DataLoader,
description="Retrieves news and news-like content from Minimap.ai content search platform.",
),

ToolName.Search_File: ManagedTool(
name=ToolName.Search_File,
implementation=SearchFileTool,
Expand Down Expand Up @@ -143,7 +162,7 @@ class ToolName(StrEnum):
def get_available_tools() -> dict[ToolName, dict]:
langchain_tools = [ToolName.Python_Interpreter, ToolName.Tavily_Internet_Search]
use_langchain_tools = bool(
strtobool(os.getenv("USE_EXPERIMENTAL_LANGCHAIN", "False"))
strtobool(os.getenv("USE_EXPERIMENTAL_LANGCHAIN", "True"))
)
use_community_tools = bool(strtobool(os.getenv("USE_COMMUNITY_FEATURES", "False")))

Expand Down
2 changes: 2 additions & 0 deletions src/backend/tools/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from backend.tools.lang_chain import LangChainVectorDBRetriever, LangChainWikiRetriever
from backend.tools.python_interpreter import PythonInterpreter
from backend.tools.tavily import TavilyInternetSearch
from backend.tools.minimap import LangChainMinimapRetriever

__all__ = [
"Calculator",
Expand All @@ -12,4 +13,5 @@
"TavilyInternetSearch",
"ReadFileTool",
"SearchFileTool",
"LangChainMinimapRetriever"
]
143 changes: 143 additions & 0 deletions src/backend/tools/minimap.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,143 @@
import os
from typing import Any, Dict, List

# from langchain.text_splitter import CharacterTextSplitter
# from langchain_cohere import CohereEmbeddings
# from langchain_community.document_loaders import PyPDFLoader
# from langchain_community.retrievers import WikipediaRetriever
# from langchain_community.vectorstores import Chroma
# from langchain import LangChain, Tool

from typing import Optional

from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.pydantic_v1 import Field
from langchain_core.tools import BaseTool

import json
import logging
import time
import requests
from typing import Any, Dict, Iterator, List

from langchain_core.documents import Document
from langchain_core.pydantic_v1 import BaseModel, root_validator

from urllib.parse import urljoin
from backend.tools.base import BaseTool as CohereBaseTool

"""
Plug in your lang chain retrieval implementation here.
We have an example flows with wikipedia and vector DBs.
More details: https://python.langchain.com/docs/integrations/retrievers
"""

logging = logging.getLogger(__name__)

class MinimapAPIWrapper(BaseModel):
"""
Wrapper around Minimap.ai API.
This wrapper will use the Minimap API to conduct searches and fetch
document summaries. By default, it will return the document summaries
of the top-k results of an input search.
Parameters:
top_k_results: number of the top-scored document
MAX_QUERY_LENGTH: maximum length of the query.
Default is 300 characters.
doc_content_chars_max: maximum length of the document content.
Content will be truncated if it exceeds this length.
Default is 2000 characters.
max_retry: maximum number of retries for a request. Default is 5.
sleep_time: time to wait between retries.
Default is 0.2 seconds.
email: email address to be used for the PubMed API.
"""

parse: Any #: :meta private:

base_url: str = "http://host.docker.internal:8081"
search_endpoint: str = urljoin(base_url, "/api/v0/platform/elastic_search")

max_retry: int = 5

# Default values for the parameters
top_k_results: int = 15
MAX_QUERY_LENGTH: int = 300
doc_content_chars_max: int = 2000


def run(self, query: str) -> str:
"""
Run search queries against the Minimap search PI
Returns a list of documents, each with a title, summary, and id.
"""
try:
# Create teh url params
query_params = {
'query': query,
}

response = requests.get(self.search_endpoint, params=query_params)


if response.status_code != 200:
return f"Minimap API returned status code {response.status_code}"

response_json = response.json()

results = response_json.get("results", [])

# limit the number of results to top_k_results
results = results[:self.top_k_results]

# strip out the `id` field from the results
# results = [{"text": result["title"], 'url': result['url']} for result in results]

print(results)
return results

except Exception as ex:
return f"PubMed exception: {ex}"


class MinimapQueryRun(BaseTool):
"""Tool that searches the Minimap.ai API."""

name: str = "pub_med"
description: str = (
"A wrapper around Minimap.ai. "
"Useful for searching news articles across a wide range of topics."
"from sports, politics, sciences, life style, and all sorts of news and news-like content."
"Input should be a search query."
)

api_wrapper: MinimapAPIWrapper = Field(default_factory=MinimapAPIWrapper)

def _run(
self,
query: str,
run_manager: Optional[CallbackManagerForToolRun] = None,
) -> str:
"""Use the PubMed tool."""
return self.api_wrapper.run(query)

class LangChainMinimapRetriever(CohereBaseTool):
"""
This class retrieves documents from Wikipedia using the langchain package.
This requires wikipedia package to be installed.
"""

def __init__(self):
self.client = MinimapAPIWrapper()

@classmethod
def is_available(cls) -> bool:
return True

def call(self, parameters: dict, **kwargs: Any) -> List[Dict[str, Any]]:
query = parameters.get("query", "")
results = self.client.run(query)
return results # [{"text": result} for result in results]

0 comments on commit 841355d

Please sign in to comment.