alita-sdk 0.3.314__py3-none-any.whl → 0.3.315__py3-none-any.whl

alita_sdk/tools/browser/__init__.py

@@ -12,6 +12,8 @@ from ..utils import get_max_toolkit_length, clean_string, TOOLKIT_SPLITTER
 from ...configurations.browser import BrowserConfiguration
 from logging import getLogger
 
+from ...configurations.pgvector import PgVectorConfiguration
+
 logger = getLogger(__name__)
 
 name = "browser"
@@ -21,6 +23,8 @@ def get_tools(tool):
     return BrowserToolkit().get_toolkit(
         selected_tools=tool['settings'].get('selected_tools', []),
         browser_configuration=tool['settings']['browser_configuration'],
+        pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
+        embedding_model=tool['settings'].get('embedding_model'),
         toolkit_name=tool.get('toolkit_name', '')
     ).get_tools()
 
@@ -51,8 +55,21 @@ class BrowserToolkit(BaseToolkit):
 
         return create_model(
             name,
-            __config__=ConfigDict(json_schema_extra={'metadata': {"label": "Browser", "icon_url": None, "categories": ["testing"], "extra_categories": ["web scraping", "search", "crawler"]}}),
-            browser_configuration=(BrowserConfiguration, Field(description="Browser Configuration", json_schema_extra={'configuration_types': ['browser']})),
+            __config__=ConfigDict(json_schema_extra={'metadata': {"label": "Browser", "icon_url": None,
+                                                                  "categories": ["testing"],
+                                                                  "extra_categories": [
+                                                                      "web scraping", "search", "crawler"
+                                                                  ]}}),
+            browser_configuration=(Optional[BrowserConfiguration],
+                                   Field(description="Browser Configuration (required for tools and `google`)",
+                                         default=None, json_schema_extra={'configuration_types': ['browser']})),
+            pgvector_configuration=(Optional[PgVectorConfiguration],
+                                    Field(description="PgVector configuration (required for tools `multi_url_crawler`)",
+                                          default=None, json_schema_extra={'configuration_types': ['pgvector']})),
+            embedding_model=(Optional[str],
+                             Field(default=None,
+                                   description="Embedding configuration (required for tools `multi_url_crawler`)",
+                                   json_schema_extra={'configuration_model': 'embedding'})),
             selected_tools=(List[Literal[tuple(selected_tools)]],
                             Field(default=[], json_schema_extra={'args_schemas': selected_tools})),
             __validators__={
@@ -65,9 +82,15 @@ class BrowserToolkit(BaseToolkit):
         if selected_tools is None:
             selected_tools = []
 
-        wrapper_payload = {
+        wrapper_payload_google = {
             **kwargs,
             **kwargs.get('browser_configuration', {}),
+            **kwargs.get('pgvector_configuration', {}),
+        }
+
+        wrapper_payload_rag_based = {
+            **kwargs,
+            **kwargs.get('pgvector_configuration', {}),
         }
 
         tools = []
@@ -85,7 +108,7 @@ class BrowserToolkit(BaseToolkit):
             if tool == 'single_url_crawler':
                 tool_entry = SingleURLCrawler()
             elif tool == 'multi_url_crawler':
-                tool_entry = MultiURLCrawler()
+                tool_entry = MultiURLCrawler(**wrapper_payload_rag_based)
             elif tool == 'get_html_content':
                 tool_entry = GetHTMLContent()
             elif tool == 'get_pdf_content':
@@ -93,7 +116,7 @@ class BrowserToolkit(BaseToolkit):
             elif tool == 'google':
                 try:
                     google_api_wrapper = GoogleSearchAPIWrapper(
-                        **wrapper_payload
+                        **wrapper_payload_google
                     )
                     tool_entry = GoogleSearchResults(api_wrapper=google_api_wrapper)
                     # rename the tool to avoid conflicts

alita_sdk/tools/browser/crawler.py

@@ -27,13 +27,16 @@ class MultiURLCrawler(BaseTool):
     max_response_size: int = 3000
     name: str = "multi_url_crawler"
     description: str = "Crawls multiple URLs and returns the content related to query"
+    embedding_model: str = None
+    connection_string: str = None
     args_schema: Type[BaseModel] = create_model("MultiURLCrawlerModel",
                                                 query=(str, Field(description="Query text to search pages")),
                                                 urls=(list[str], Field(description="list of URLs to search like ['url1', 'url2']")))
 
     def _run(self, query: str, urls: list[str], run_manager=None):
         urls = [url.strip() for url in urls]
-        return webRag(urls, self.max_response_size, query)
+        return webRag(urls=urls, max_response_size=self.max_response_size, query=query,
+                      connection_string=self.connection_string, embedding_model=self.embedding_model)
 
 
 class GetHTMLContent(BaseTool):

alita_sdk/tools/browser/utils.py

@@ -6,9 +6,9 @@ from langchain.text_splitter import CharacterTextSplitter
 import fitz
 
 try:
-    from langchain_chroma import Chroma
+    from langchain_postgres import PGVector
 except ImportError:
-    Chroma = None
+    PGVector = None
 
 from langchain_community.embeddings.sentence_transformer import (
     SentenceTransformerEmbeddings,
@@ -32,13 +32,22 @@ def get_page(urls, html_only=False):
     return docs_transformed
 
 
-def webRag(urls, max_response_size, query):
-    if Chroma is None:
-        return "Chroma is not initialized. Web rag is not available."
+def webRag(urls, max_response_size, query, connection_string=None, embedding_model=None):
+    if PGVector is None:
+        return "PGVector is not initialized. Web rag is not available."
+
+    if not connection_string or not embedding_model:
+        return "Connection string or embedding model is missing. Web rag is not available."
     text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
     docs = text_splitter.split_documents(get_page(urls))
-    embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
-    db = Chroma.from_documents(docs, embedding_function)
+    embedding_function = SentenceTransformerEmbeddings(model_name=embedding_model)
+    db = PGVector.from_documents(
+        documents=docs,
+        embedding=embedding_function,
+        collection_name="web_rag",
+        pre_delete_collection=True,
+        connection=connection_string
+    )
     docs = db.search(query, "mmr", k=10)
     text = ""
     for doc in docs:

{alita_sdk-0.3.314.dist-info → alita_sdk-0.3.315.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.314
+Version: 0.3.315
 Summary: SDK for building langchain agents using resources from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedj27@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0

{alita_sdk-0.3.314.dist-info → alita_sdk-0.3.315.dist-info}/RECORD

@@ -162,11 +162,11 @@ alita_sdk/tools/bitbucket/__init__.py,sha256=2VAY45Jij5dHkz6UGTmsEmOcLeJMWmcX-Wr
 alita_sdk/tools/bitbucket/api_wrapper.py,sha256=xHrluV2aCckOK_lGd42fFz1c-pyuZAnC-A_v1SKiM5g,20006
 alita_sdk/tools/bitbucket/bitbucket_constants.py,sha256=UsbhQ1iEvrKoxceTFPWTYhaXS1zSxbmjs1TwY0-P4gw,462
 alita_sdk/tools/bitbucket/cloud_api_wrapper.py,sha256=QHdud-d3xcz3mOP3xb1Htk1sv9QFg7bTm1szdN_zohQ,15517
-alita_sdk/tools/browser/__init__.py,sha256=ypYaShyDzjKDgErANsaMohRv0C_fQaVANX0As1L3amQ,5316
-alita_sdk/tools/browser/crawler.py,sha256=jhE35dU94eQLURSM-D50tspOqEMsiGzMDbYNqNSR2mU,2279
+alita_sdk/tools/browser/__init__.py,sha256=faLmuMt2CzCPMmxquGcdV-TGAbKxHi7sTQyuK0VKYNs,6760
+alita_sdk/tools/browser/crawler.py,sha256=Ah0tyF7lKGJIlxMY4MXEQmuDehaB_I-FmECxG27DnPw,2476
 alita_sdk/tools/browser/duck_duck_go_search.py,sha256=iKws923v34o-ySXohJw-8xTDBWlj3fMsnzC_ZRuPugE,2002
 alita_sdk/tools/browser/google_search_rag.py,sha256=QVHFbVwymiJGuno_HLSJOK1c_MpgMdBSTYQKf6fLRk8,1838
-alita_sdk/tools/browser/utils.py,sha256=4k3YM_f1Kqlhjz9vt2pNsGkvCjhy-EmY3nvcwdFCsLA,2501
+alita_sdk/tools/browser/utils.py,sha256=J4-ZSb5TeCJnYJTsPKUOyiOC_vfTye3QtZA-T_AYEoA,2853
 alita_sdk/tools/browser/wiki.py,sha256=Qh3HBFd4dkS2VavXbFJOm4b8SjVSIe5xSD7CY1vEkKE,1126
 alita_sdk/tools/carrier/__init__.py,sha256=Ove5wAXBxyLS5F5ZxgydV2xKZJIR3OoMB5fMkn8jNUc,4296
 alita_sdk/tools/carrier/api_wrapper.py,sha256=tP7oR_U0HX1rxqat0Jkz6oh3RB9BEr1ESKQ9J8OWDcE,9093
@@ -349,8 +349,8 @@ alita_sdk/tools/zephyr_scale/api_wrapper.py,sha256=A6CUEKjENt3mZlPU9lai88WV9esCD
 alita_sdk/tools/zephyr_squad/__init__.py,sha256=0ne8XLJEQSLOWfzd2HdnqOYmQlUliKHbBED5kW_Vias,2895
 alita_sdk/tools/zephyr_squad/api_wrapper.py,sha256=kmw_xol8YIYFplBLWTqP_VKPRhL_1ItDD0_vXTe_UuI,14906
 alita_sdk/tools/zephyr_squad/zephyr_squad_cloud_client.py,sha256=R371waHsms4sllHCbijKYs90C-9Yu0sSR3N4SUfQOgU,5066
-alita_sdk-0.3.314.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-alita_sdk-0.3.314.dist-info/METADATA,sha256=zmL9d8KWnZmuKvbhUXrbLn4HfljL2bvYyilxieU63Ys,18897
-alita_sdk-0.3.314.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-alita_sdk-0.3.314.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
-alita_sdk-0.3.314.dist-info/RECORD,,
+alita_sdk-0.3.315.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+alita_sdk-0.3.315.dist-info/METADATA,sha256=iwfMJ3l6wJ6OAHVEJWZ1xIUsQgfKlcgUzxdqUvS4POA,18897
+alita_sdk-0.3.315.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+alita_sdk-0.3.315.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
+alita_sdk-0.3.315.dist-info/RECORD,,