Merge pull request #1 from joaomdmoura/gui/rag-tool

First take on a rag tool

README.md

@@ -0,0 +1,101 @@
+## Getting started
+
+When setting up agents you can provide tools for them to use. Here you will find ready-to-use tools as well as simple helpers for you to create your own tools.
+
+In order to create a new tool, you have to pick one of the available strategies.
+
+### Subclassing `BaseTool`
+
+```python
+class MyTool(BaseTool):
+    name: str = "Knowledge base"
+    description: str = "A knowledge base with all the requirements for the project."
+
+    def _run(self, question) -> str:
+        return (
+            tbl.search(embed_func([question])[0]).limit(3).to_pandas()["text"].tolist()
+        )
+```
+
+As you can see, all you need to do is to create a new class that inherits from `BaseTool`, define `name` and `description` fields, as well as implement the `_run` method.
+
+### Create tool from a function or lambda
+
+```python
+my_tool = Tool(
+    name="Knowledge base",
+    description="A knowledge base with all the requirements for the project.",
+    func=lambda question: tbl.search(embed_func([question])[0])
+    .limit(3)
+    .to_pandas()["text"]
+    .tolist(),
+)
+```
+
+Here's it's a bit simpler, as you don't have to subclass. Simply create a `Tool` object with the three required fields and you are good to go.
+
+### Use the `tool` decorator.
+
+```python
+@tool("Knowledge base")
+def my_tool(question: str) -> str:
+    """A knowledge base with all the requirements for the project."""
+    return tbl.search(embed_func([question])[0]).limit(3).to_pandas()["text"].tolist()
+```
+
+By using the decorator you can easily wrap simple functions as tools. If you don't provide a name, the function name is going to be used. However, the docstring is required.
+
+If you are using a linter you may see issues when passing your decorated tool in `tools` parameters that expect a list of `BaseTool`. If that's the case, you can use the `as_tool` helper.
+
+
+## Contribution
+
+This repo is open-source and we welcome contributions. If you're looking to contribute, please:
+
+- Fork the repository.
+- Create a new branch for your feature.
+- Add your feature or improvement.
+- Send a pull request.
+- We appreciate your input!
+
+### Installing Dependencies
+
+```bash
+poetry install
+```
+
+### Virtual Env
+
+```bash
+poetry shell
+```
+
+### Pre-commit hooks
+
+```bash
+pre-commit install
+```
+
+### Running Tests
+
+```bash
+poetry run pytest
+```
+
+### Running static type checks
+
+```bash
+poetry run pyright
+```
+
+### Packaging
+
+```bash
+poetry build
+```
+
+### Installing Locally
+
+```bash
+pip install dist/*.tar.gz
+```

src/crewai_tools/__init__.py

@@ -0,0 +1 @@
+from .base_tool import BaseTool, Tool, as_tool, tool

src/crewai_tools/adapters/embedchain_adapter.py

@@ -0,0 +1,14 @@
+from embedchain import App
+
+from crewai_tools.rag_tool import Adapter
+
+
+class EmbedchainAdapter(Adapter):
+    embedchain_app: App
+    dry_run: bool = False
+
+    def query(self, question: str) -> str:
+        result = self.embedchain_app.query(question, dry_run=self.dry_run)
+        if result is list:
+            return "\n".join(result)
+        return str(result)

src/crewai_tools/adapters/lancedb_adapter.py

@@ -0,0 +1,49 @@
+from pathlib import Path
+from typing import Any, Callable
+
+from lancedb import DBConnection as LanceDBConnection
+from lancedb import connect as lancedb_connect
+from lancedb.table import Table as LanceDBTable
+from openai import Client as OpenAIClient
+from pydantic import Field, PrivateAttr
+
+from crewai_tools.rag_tool import Adapter
+
+
+def _default_embedding_function():
+    client = OpenAIClient()
+
+    def _embedding_function(input):
+        rs = client.embeddings.create(input=input, model="text-embedding-ada-002")
+        return [record.embedding for record in rs.data]
+
+    return _embedding_function
+
+
+class LanceDBAdapter(Adapter):
+    uri: str | Path
+    table_name: str
+    embedding_function: Callable = Field(default_factory=_default_embedding_function)
+    top_k: int = 3
+    vector_column_name: str = "vector"
+    text_column_name: str = "text"
+
+    _db: LanceDBConnection = PrivateAttr()
+    _table: LanceDBTable = PrivateAttr()
+
+    def model_post_init(self, __context: Any) -> None:
+        self._db = lancedb_connect(self.uri)
+        self._table = self._db.open_table(self.table_name)
+
+        return super().model_post_init(__context)
+
+    def query(self, question: str) -> str:
+        query = self.embedding_function([question])[0]
+        results = (
+            self._table.search(query, vector_column_name=self.vector_column_name)
+            .limit(self.top_k)
+            .select([self.text_column_name])
+            .to_list()
+        )
+        values = [result[self.text_column_name] for result in results]
+        return "\n".join(values)

src/crewai_tools/base_tool.py

@@ -0,0 +1,82 @@
+from abc import ABC, abstractmethod
+from typing import Any, Callable, cast
+
+from langchain.agents import tools as langchain_tools
+from pydantic import BaseModel
+
+
+class BaseTool(BaseModel, ABC):
+    name: str
+    """The unique name of the tool that clearly communicates its purpose."""
+    description: str
+    """Used to tell the model how/when/why to use the tool."""
+
+    def run(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        return self._run(*args, **kwargs)
+
+    @abstractmethod
+    def _run(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        """Here goes the actual implementation of the tool."""
+
+    def to_langchain(self) -> langchain_tools.Tool:
+        return langchain_tools.Tool(
+            name=self.name,
+            description=self.description,
+            func=self._run,
+        )
+
+
+class Tool(BaseTool):
+    func: Callable
+    """The function that will be executed when the tool is called."""
+
+    def _run(self, *args: Any, **kwargs: Any) -> Any:
+        return self.func(*args, **kwargs)
+
+
+def to_langchain(
+    tools: list[BaseTool | langchain_tools.BaseTool],
+) -> list[langchain_tools.BaseTool]:
+    return [t.to_langchain() if isinstance(t, BaseTool) else t for t in tools]
+
+
+def tool(*args):
+    """
+    Decorator to create a tool from a function.
+    """
+
+    def _make_with_name(tool_name: str) -> Callable:
+        def _make_tool(f: Callable) -> BaseTool:
+            if f.__doc__ is None:
+                raise ValueError("Function must have a docstring")
+
+            return Tool(
+                name=tool_name,
+                description=f.__doc__,
+                func=f,
+            )
+
+        return _make_tool
+
+    if len(args) == 1 and callable(args[0]):
+        return _make_with_name(args[0].__name__)(args[0])
+    if len(args) == 1 and isinstance(args[0], str):
+        return _make_with_name(args[0])
+    raise ValueError("Invalid arguments")
+
+
+def as_tool(f: Any) -> BaseTool:
+    """
+    Useful for when you create a tool using the @tool decorator and want to use it as a BaseTool.
+    It is a BaseTool, but type inference doesn't know that.
+    """
+    assert isinstance(f, BaseTool)
+    return cast(BaseTool, f)

src/crewai_tools/rag_tool.py

@@ -0,0 +1,74 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+from crewai_tools.base_tool import BaseTool
+
+
+class Adapter(BaseModel, ABC):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    @abstractmethod
+    def query(self, question: str) -> str:
+        """Query the knowledge base with a question and return the answer."""
+
+
+class RagTool(BaseTool):
+    name: str = "Knowledge base"
+    description: str = "A knowledge base that can be used to answer questions."
+    adapter: Adapter
+
+    def _run(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        return self.adapter.query(args[0])
+
+    def from_file(self, file_path: str):
+        from embedchain import App
+        from embedchain.models.data_type import DataType
+
+        from crewai_tools.adapters.embedchain_adapter import EmbedchainAdapter
+
+        app = App()
+        app.add(file_path, data_type=DataType.TEXT_FILE)
+
+        adapter = EmbedchainAdapter(embedchain_app=app)
+        return RagTool(adapter=adapter)
+
+    def from_directory(self, directory_path: str):
+        from embedchain import App
+        from embedchain.loaders.directory_loader import DirectoryLoader
+
+        from crewai_tools.adapters.embedchain_adapter import EmbedchainAdapter
+
+        loader = DirectoryLoader(config=dict(recursive=True))
+
+        app = App()
+        app.add(directory_path, loader=loader)
+
+        adapter = EmbedchainAdapter(embedchain_app=app)
+        return RagTool(adapter=adapter)
+
+    def from_web_page(self, url: str):
+        from embedchain import App
+        from embedchain.models.data_type import DataType
+
+        from crewai_tools.adapters.embedchain_adapter import EmbedchainAdapter
+
+        app = App()
+        app.add(url, data_type=DataType.WEB_PAGE)
+
+        adapter = EmbedchainAdapter(embedchain_app=app)
+        return RagTool(adapter=adapter)
+
+    def from_embedchain(self, config_path: str):
+        from embedchain import App
+
+        from crewai_tools.adapters.embedchain_adapter import EmbedchainAdapter
+
+        app = App.from_config(config_path=config_path)
+        adapter = EmbedchainAdapter(embedchain_app=app)
+        return RagTool(adapter=adapter)

tests/adapters/embedchain_adapter_test.py

@@ -0,0 +1,67 @@
+from typing import Callable
+
+from chromadb import Documents, EmbeddingFunction, Embeddings
+from embedchain import App
+from embedchain.config import AppConfig, ChromaDbConfig
+from embedchain.embedder.base import BaseEmbedder
+from embedchain.vectordb.chroma import ChromaDB
+
+from crewai_tools.adapters.embedchain_adapter import EmbedchainAdapter
+
+
+class MockEmbeddingFunction(EmbeddingFunction):
+    fn: Callable
+
+    def __init__(self, embedding_fn: Callable):
+        self.fn = embedding_fn
+
+    def __call__(self, input: Documents) -> Embeddings:
+        return self.fn(input)
+
+
+def test_embedchain_adapter(helpers):
+    embedding_function = MockEmbeddingFunction(
+        embedding_fn=helpers.get_embedding_function()
+    )
+    embedder = BaseEmbedder()
+    embedder.set_embedding_fn(embedding_function)  # type: ignore
+
+    db = ChromaDB(
+        config=ChromaDbConfig(
+            dir="tests/data/chromadb",
+            collection_name="requirements",
+        )
+    )
+
+    app = App(
+        config=AppConfig(
+            id="test",
+        ),
+        db=db,
+        embedding_model=embedder,
+    )
+
+    adapter = EmbedchainAdapter(
+        dry_run=True,
+        embedchain_app=app,
+    )
+
+    assert (
+        adapter.query("What are the requirements for the task?")
+        == """
+  Use the following pieces of context to answer the query at the end.
+  If you don't know the answer, just say that you don't know, don't try to make up an answer.
+
+  Technical requirements
+
+The system should be able to process 1000 transactions per second. The code must be written in Ruby. | Problem
+
+Currently, we are not able to find out palindromes in a given string. We need a solution to this problem. | Solution
+
+We need a function that takes a string as input and returns true if the string is a palindrome, otherwise false.
+
+  Query: What are the requirements for the task?
+
+  Helpful Answer:
+"""
+    )

tests/adapters/lancedb_adapter_test.py

@@ -0,0 +1,22 @@
+from crewai_tools.adapters.lancedb_adapter import LanceDBAdapter
+
+
+def test_lancedb_adapter(helpers):
+    adapter = LanceDBAdapter(
+        uri="tests/data/lancedb",
+        table_name="requirements",
+        embedding_function=helpers.get_embedding_function(),
+        top_k=2,
+        vector_column_name="vector",
+        text_column_name="text",
+    )
+
+    assert (
+        adapter.query("What are the requirements for the task?")
+        == """Technical requirements
+
+The system should be able to process 1000 transactions per second. The code must be written in Ruby.
+Problem
+
+Currently, we are not able to find out palindromes in a given string. We need a solution to this problem."""
+    )

tests/conftest.py

@@ -0,0 +1,21 @@
+from typing import Callable
+
+import pytest
+
+
+class Helpers:
+    @staticmethod
+    def get_embedding_function() -> Callable:
+        def _func(input):
+            assert input == ["What are the requirements for the task?"]
+            with open("tests/data/embedding.txt", "r") as file:
+                content = file.read()
+                numbers = content.split(",")
+                return [[float(number) for number in numbers]]
+
+        return _func
+
+
+@pytest.fixture
+def helpers():
+    return Helpers

tests/data/embedding.txt

@@ -0,0 +1 @@
+0.025692760944366455, -0.004980836994946003, 0.008317121304571629, -0.013715836219489574, -0.02352249063551426, 0.024290846660733223, -0.02944018319249153, -0.00994145404547453, -0.009779694490134716, -0.027930431067943573, 0.015367128886282444, -0.008276681415736675, -0.014302213676273823, -0.009388776496052742, 0.00832386128604412, 0.005149336066097021, 0.027593431994318962, -0.010999629274010658, -0.0037170927971601486, -0.009934714064002037, -0.011660145595669746, -0.0008045837748795748, -0.003956361673772335, 0.006497330032289028, -0.011350107379257679, 0.00026391190476715565, 0.02470872551202774, -0.018925832584500313, 0.027418192476034164, -0.01903367228806019, 0.01558280736207962, -0.024061689153313637, -0.02167573943734169, -0.016270285472273827, -0.01949198916554451, 0.0019411110552027822, 0.0031981151551008224, -0.020031187683343887, 0.02670375630259514, 0.006214251276105642, 0.010089733637869358, 0.028092190623283386, 0.01414045412093401, 0.0019175211200490594, -0.01567716710269451, -0.0036530629731714725, -0.021567899733781815, -0.019896388053894043, -0.022781094536185265, -0.004411309491842985, 0.017240840941667557, 0.0107906898483634, -0.03022201918065548, -0.02000422775745392, -0.01438309345394373, 0.004252920392900705, 0.0053582750260829926, 0.0011028273729607463, -0.002139940159395337, -0.007750964257866144, -0.002953791292384267, 0.009193317033350468, -0.0018888761987909675, -0.013452977873384953, -0.0053515350446105, 7.161216672102455e-06, -0.027714751660823822, 0.017106041312217712, -0.0243986863642931, -0.013816935941576958, 0.011727545410394669, 0.008579980581998825, 0.017362158745527267, 0.026488076895475388, 0.03019505925476551, -0.024317806586623192, 0.019235869869589806, -0.011181607842445374, -0.007831843569874763, 0.0052234758622944355, 0.010689590126276016, -0.004522519186139107, -0.02803827077150345, 0.03092297725379467, 0.004411309491842985, -0.0002287376846652478, 0.020597344264388084, 0.015218849293887615, -0.0026269028894603252, -0.011997144669294357, -0.000433464243542403, 0.01713300123810768, 0.025962360203266144, 0.0021685848478227854, 0.014019135385751724, 0.01903367228806019, -0.003508153837174177, 0.007326345890760422, -0.015286249108612537, -0.021244380623102188, -0.02191837877035141, 0.006305240560323, -0.038336943835020065, -0.007050007116049528, -0.0011550621129572392, -0.01621636562049389, 0.00021378337987698615, -0.0058064828626811504, 0.011936484836041927, 0.0009318006923422217, -0.02216101810336113, 0.00864063948392868, -0.002648807829245925, -0.01830575428903103, -0.0008745109662413597, 0.0035452235024422407, 0.0331067256629467, -0.009227016940712929, 0.0033312295563519, -0.01769915781915188, 0.033942483365535736, 0.0009014708339236677, 0.02021990530192852, -0.006891618017107248, 0.029035786166787148, -0.016782522201538086, 0.01173428539186716, -0.010817649774253368, -0.010837869718670845, -0.020597344264388084, 0.012765500694513321, 0.018831472843885422, 0.007211766671389341, 0.025301842018961906, -0.012408282607793808, 0.008842838928103447, -0.017173441126942635, 0.016782522201538086, -0.01964026875793934, -0.0016723547596484423, -0.006847808137536049, 0.02189141884446144, 0.00038122947444207966, -0.001764186774380505, 0.0025746680330485106, 0.0011213623220100999, 0.01900671236217022, 0.004953877069056034, -0.002926831366494298, -0.021244380623102188, -0.000591011019423604, -0.0057357135228812695, 0.015407568775117397, -0.006274911109358072, 0.002881336724385619, 0.007225246634334326, -0.025113124400377274, 0.013365358114242554, -0.0008104812586680055, -0.015097529627382755, 0.010844609700143337, 0.006082821637392044, 0.005641353782266378, -0.007258946541696787, 0.016148965805768967, 0.03019505925476551, 0.021567899733781815, -0.0033952591475099325, -0.006325460504740477, 0.00723872659727931, 0.015528888441622257, -0.0021972297690808773, -0.02944018319249153, 0.00900459848344326, 0.021837498992681503, 0.01800919696688652, 0.011788205243647099, -0.012246523052453995, -0.013668656349182129, -0.002812251914292574, -0.008344081230461597, 0.020233385264873505, 0.004542739130556583, 0.00920005701482296, -0.02398080937564373, -0.0023421391379088163, -0.005742453504353762, -0.009098958224058151, 0.00953705608844757, -0.008317121304571629, 0.025059204548597336, -0.010238012298941612, 0.008916978724300861, -0.009948194026947021, -0.6888787150382996, -0.009193317033350468, 0.03343024477362633, -0.02873922698199749, -0.00014174997340887785, 0.023697730153799057, -0.0012350992765277624, 0.008654119446873665, -0.00398669159039855, 0.0304646585136652, 0.005186405964195728, -0.0051931459456682205, 0.02055690437555313, -0.021015223115682602, -0.014827931299805641, -0.02871226705610752, 0.00853280071169138, -4.6784862206550315e-05, -0.0058098528534173965, 0.02152745984494686, -0.03210921213030815, 0.010029073804616928, 0.0010969298891723156, -0.0033093246165663004, 0.006588319316506386, 0.0008395473705604672, 0.01066937018185854, -0.0343199223279953, 0.007326345890760422, Line truncated

tests/data/lancedb/requirements.lance/_transactions/0-d2c46569-d173-4b3f-b589-f8f00eddc371.txn

@@ -0,0 +1 @@
+$d2c46569-d173-4b3f-b589-f8f00eddc371²Vtext ÿÿÿÿÿÿÿÿÿ
*string0
85vector
 ÿÿÿÿÿÿÿÿÿ
*fixed_size_list:float:15360
8

tests/rag_tool_test.py

@@ -0,0 +1,21 @@
+from crewai_tools.rag_tool import Adapter, RagTool
+
+
+class MockAdapter(Adapter):
+    answer: str
+
+    def query(self, question: str) -> str:
+        return self.answer
+
+
+def test_rag_tool():
+    adapter = MockAdapter(answer="42")
+    rag_tool = RagTool(adapter=adapter)
+
+    assert rag_tool.name == "Knowledge base"
+    assert (
+        rag_tool.description == "A knowledge base that can be used to answer questions."
+    )
+    assert (
+        rag_tool.run("What is the answer to life, the universe and everything?") == "42"
+    )