alita-sdk 0.3.205__py3-none-any.whl → 0.3.206__py3-none-any.whl

alita_sdk/runtime/clients/client.py

@@ -10,6 +10,7 @@ from langchain_core.messages import (
 )
 from langchain_core.tools import ToolException
 from langgraph.store.base import BaseStore
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 
 from ..langchain.assistant import Assistant as LangChainAssistant
 # from ..llamaindex.assistant import Assistant as LLamaAssistant
@@ -37,6 +38,7 @@ class AlitaClient:
 
         self.base_url = base_url.rstrip('/')
         self.api_path = '/api/v1'
+        self.llm_path = '/llm/v1'
         self.project_id = project_id
         self.auth_token = auth_token
         self.headers = {
@@ -152,6 +154,35 @@ class AlitaClient:
             return resp.json()
         return []
 
+    def get_llm(self, model_name: str, model_config: dict) -> ChatOpenAI:
+        """
+        Get a ChatOpenAI model instance based on the model name and configuration.
+        
+        Args:
+            model_name: Name of the model to retrieve
+            model_config: Configuration parameters for the model
+            
+        Returns:
+            An instance of ChatOpenAI configured with the provided parameters.
+        """
+        if not model_name:
+            raise ValueError("Model name must be provided")
+        
+        logger.info(f"Creating ChatOpenAI model: {model_name} with config: {model_config}")
+        
+        return ChatOpenAI(
+            base_url=f"{self.base_url}{self.llm_path}",
+            model=model_name,
+            api_key=self.auth_token,
+            stream_usage=model_config.get("stream_usage", True),
+            max_tokens=model_config.get("max_tokens", None),
+            top_p=model_config.get("top_p"),
+            temperature=model_config.get("temperature"),
+            max_retries=model_config.get("max_retries", 3),
+            seed=model_config.get("seed", None),
+        )
+         
+    
     def get_app_version_details(self, application_id: int, application_version_id: int) -> dict:
         url = f"{self.application_versions}/{application_id}/{application_version_id}"
         if self.configurations:
@@ -177,11 +208,12 @@ class AlitaClient:
         logger.info(f"Unsecret response: {data}")
         return data.get('value', None)
 
-    def application(self, client: Any, application_id: int, application_version_id: int,
+    def application(self, application_id: int, application_version_id: int,
                     tools: Optional[list] = None, chat_history: Optional[List[Any]] = None,
                     app_type=None, memory=None, runtime='langchain',
                     application_variables: Optional[dict] = None,
-                    version_details: Optional[dict] = None, store: Optional[BaseStore] = None):
+                    version_details: Optional[dict] = None, store: Optional[BaseStore] = None,
+                    llm: Optional[ChatOpenAI] = None):
         if tools is None:
             tools = []
         if chat_history is None:
@@ -200,7 +232,15 @@ class AlitaClient:
             for var in data.get('variables', {}):
                 if var['name'] in application_variables:
                     var.update(application_variables[var['name']])
-
+        if llm is None:
+            llm = self.get_llm(
+                model_name=data['llm_settings']['model_name'],
+                model_config = {
+                    "max_tokens": data['llm_settings']['max_tokens'],
+                    "top_p": data['llm_settings']['top_p'],
+                    "temperature": data['llm_settings']['temperature']
+                }
+            )
         if not app_type:
             app_type = data.get("agent_type", "react")
         if app_type == "alita":
@@ -212,10 +252,10 @@ class AlitaClient:
         elif app_type == 'autogen':
             app_type = "openai"
         if runtime == 'nonrunnable':
-            return LangChainAssistant(self, data, client, chat_history, app_type,
+            return LangChainAssistant(self, data, llm, chat_history, app_type,
                                       tools=tools, memory=memory, store=store)
         if runtime == 'langchain':
-            return LangChainAssistant(self, data, client,
+            return LangChainAssistant(self, data, llm,
                                       chat_history, app_type,
                                       tools=tools, memory=memory, store=store).runnable()
         elif runtime == 'llama':

alita_sdk/runtime/langchain/assistant.py

@@ -36,29 +36,30 @@ class Assistant:
 
         logger.debug("Data for agent creation: %s", data)
         logger.info("App type: %s", app_type)
-
+        
+        self.client = client
         # For predict agents, use the client as-is since it's already configured
-        if app_type == "predict":
-            self.client = client
-        else:
-            # For other agent types, configure client from llm_settings
-            self.client = copy(client)
-            self.client.max_tokens = data['llm_settings']['max_tokens']
-            self.client.temperature = data['llm_settings']['temperature']
-            self.client.top_p = data['llm_settings']['top_p']
-            self.client.top_k = data['llm_settings']['top_k']
-            self.client.model_name = data['llm_settings']['model_name']
-            self.client.integration_uid = data['llm_settings']['integration_uid']
+        # if app_type == "predict":
+        #     self.client = client
+        # else:
+        #     # For other agent types, configure client from llm_settings
+        #     self.client = copy(client)
+        #     self.client.max_tokens = data['llm_settings']['max_tokens']
+        #     self.client.temperature = data['llm_settings']['temperature']
+        #     self.client.top_p = data['llm_settings']['top_p']
+        #     self.client.top_k = data['llm_settings']['top_k']
+        #     self.client.model_name = data['llm_settings']['model_name']
+        #     self.client.integration_uid = data['llm_settings']['integration_uid']
 
-            model_type = data["llm_settings"]["indexer_config"]["ai_model"]
-            model_params = data["llm_settings"]["indexer_config"]["ai_model_params"]
-            #
-            target_pkg, target_name = model_type.rsplit(".", 1)
-            target_cls = getattr(
-                importlib.import_module(target_pkg),
-                target_name
-            )
-            self.client = target_cls(**model_params)
+        #     model_type = data["llm_settings"]["indexer_config"]["ai_model"]
+        #     model_params = data["llm_settings"]["indexer_config"]["ai_model_params"]
+        #     #
+        #     target_pkg, target_name = model_type.rsplit(".", 1)
+        #     target_cls = getattr(
+        #         importlib.import_module(target_pkg),
+        #         target_name
+        #     )
+        #     self.client = target_cls(**model_params)
         # validate agents compatibility: non-pipeline agents cannot have pipelines as toolkits
         if app_type not in ["pipeline", "predict"]:
             tools_to_check = data.get('tools', [])

alita_sdk/runtime/langchain/interfaces/llm_processor.py

@@ -50,9 +50,6 @@ def get_model(model_type: str, model_params: dict):
         return get_llm(model_type)(**model_params)
     if model_type == "PreloadedChatModel":
         return PreloadedChatModel(**model_params)
-    if model_type == "Alita":
-        from ...llms.alita import AlitaClient
-        return AlitaClient(**model_params)
     if model_type in chat_models:
         model = getattr(
             __import__("langchain_community.chat_models", fromlist=[model_type]),
@@ -185,7 +182,7 @@ def add_documents(vectorstore, documents):
         texts.append(document.page_content)
         for key in document.metadata:
             if isinstance(document.metadata[key], list):
-                document.metadata[key] = "; ".join(document.metadata[key])
+                document.metadata[key] = "; ".join([str(val) for val in document.metadata[key]])
             if isinstance(document.metadata[key], dict):
                 document.metadata[key] = dumps(document.metadata[key])
         metadata.append(document.metadata)

alita_sdk/runtime/toolkits/application.py

@@ -21,25 +21,20 @@ class ApplicationToolkit(BaseToolkit):
         )
     
     @classmethod
-    def get_toolkit(cls, client: Any, application_id: int, application_version_id: int, app_api_key: str,
+    def get_toolkit(cls, client: 'AlitaClient', application_id: int, application_version_id: int,
                     selected_tools: list[str] = [], store: Optional[BaseStore] = None):
-        from ..llms.alita import AlitaChatModel
         
         app_details = client.get_app_details(application_id)
         version_details = client.get_app_version_details(application_id, application_version_id)
-        settings = {
-            "deployment": client.base_url,
-            "model": version_details['llm_settings']['model_name'],
-            "api_key": app_api_key,
-            "project_id": client.project_id,
-            "integration_uid": version_details['llm_settings']['integration_uid'],
+        model_settings = {
             "max_tokens": version_details['llm_settings']['max_tokens'],
             "top_p": version_details['llm_settings']['top_p'],
-            "top_k": version_details['llm_settings']['top_k'],
             "temperature": version_details['llm_settings']['temperature'],
         }
 
-        app = client.application(AlitaChatModel(**settings), application_id, application_version_id, store=store)
+        app = client.application(application_id, application_version_id, store=store, 
+                                 llm=client.get_llm(version_details['llm_settings']['model_name'], 
+                                                    model_settings))
         return cls(tools=[Application(name=app_details.get("name"), 
                                       description=app_details.get("description"), 
                                       application=app, 

alita_sdk/runtime/toolkits/tools.py

@@ -54,7 +54,6 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                 alita_client,
                 application_id=int(tool['settings']['application_id']),
                 application_version_id=int(tool['settings']['application_version_id']),
-                app_api_key=alita_client.auth_token,
                 selected_tools=[]
             ).get_tools())
         elif tool['type'] == 'application' and tool.get('agent_type', '') == 'pipeline':

alita_sdk/runtime/tools/vectorstore.py

@@ -1,6 +1,9 @@
 import json
 import math
-from typing import Any, Optional, List, Dict, Callable
+import types
+from typing import Any, Optional, List, Dict, Callable, Generator
+
+from langchain_core.documents import Document
 from pydantic import BaseModel, model_validator, Field
 from ..langchain.tools.vector import VectorAdapter
 from langchain_core.messages import HumanMessage
@@ -8,6 +11,7 @@ from alita_sdk.tools.elitea_base import BaseToolApiWrapper
 from logging import getLogger
 
 from ..utils.logging import dispatch_custom_event
+from ..utils.utils import IndexerKeywords
 
 logger = getLogger(__name__)
 
@@ -197,25 +201,29 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             for doc_str, meta, db_id in zip(data['documents'], data['metadatas'], data['ids']):
                 doc = json.loads(doc_str)
                 doc_id = str(meta['id'])
+                dependent_docs = meta.get(IndexerKeywords.DEPENDENT_DOCS.value, [])
+                parent_id = meta.get(IndexerKeywords.PARENT.value, -1)
                 result[doc_id] = {
                     'metadata': meta,
                     'document': doc,
-                    'id': db_id
+                    'id': db_id,
+                    IndexerKeywords.DEPENDENT_DOCS.value: dependent_docs,
+                    IndexerKeywords.PARENT.value: parent_id
                 }
         except Exception as e:
             logger.error(f"Failed to get indexed data from vectorstore: {str(e)}. Continuing with empty index.")
         return result
 
-    def _reduce_duplicates(self, documents, store) -> List[Any]:
+    def _reduce_duplicates(self, documents: Generator[Document, None, None], store) -> List[Any]:
         """Remove documents already indexed in the vectorstore based on metadata 'id' and 'updated_on' fields."""
 
         self._log_data("Verification of documents to index started", tool_name="index_documents")
 
-        data = self._get_indexed_data(store)
-        indexed_ids = set(data.keys())
+        indexed_data = self._get_indexed_data(store)
+        indexed_ids = set(indexed_data.keys())
         if not indexed_ids:
             self._log_data("Vectorstore is empty, indexing all incoming documents", tool_name="index_documents")
-            return documents
+            return list(documents)
 
         final_docs = []
         docs_to_remove = []
@@ -226,13 +234,17 @@ class VectorStoreWrapper(BaseToolApiWrapper):
             if doc_id in indexed_ids:
                 # document has been indexed already, then verify `updated_on`
                 to_index_updated_on = document.metadata.get('updated_on')
-                indexed_meta = data[doc_id]['metadata']
+                indexed_meta = indexed_data[doc_id]['metadata']
                 indexed_updated_on = indexed_meta.get('updated_on')
                 if to_index_updated_on and indexed_updated_on and to_index_updated_on == indexed_updated_on:
                     # same updated_on, skip indexing
                     continue
                 # if updated_on is missing or different, we will re-index the document and remove old one
-                docs_to_remove.append(data[doc_id]['id'])
+                # parent doc removal
+                docs_to_remove.append(indexed_data[doc_id]['id'])
+                # mark dependent docs for removal
+                for dependent_doc_id in indexed_data[doc_id][IndexerKeywords.DEPENDENT_DOCS.value]:
+                    docs_to_remove.append(indexed_data[dependent_doc_id]['id'])
             else:
                 final_docs.append(document)
 
@@ -245,7 +257,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
 
         return final_docs
 
-    def index_documents(self, documents, progress_step: int = 20, clean_index: bool = True):
+    def index_documents(self, documents: Generator[Document, None, None], progress_step: int = 20, clean_index: bool = True):
         """ Index documents in the vectorstore.
 
         Args:
@@ -269,6 +281,8 @@ class VectorStoreWrapper(BaseToolApiWrapper):
                                tool_name="index_documents")
             except Exception as e:
                 logger.warning(f"Failed to clean index: {str(e)}. Continuing with re-indexing.")
+            if isinstance(documents, types.GeneratorType):
+                documents = list(documents)
         else:
             # remove duplicates based on metadata 'id' and 'updated_on' fields
             documents = self._reduce_duplicates(documents, self.vectoradapter.vectorstore)
@@ -281,7 +295,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
         # if func is provided, apply it to documents
         # used for processing of documents before indexing,
         # e.g. to avoid time-consuming operations for documents that are already indexed
-        self.process_document_func(documents) if self.process_document_func else None
+        dependent_docs_generator = self.process_document_func(documents) if self.process_document_func else []
 
         # notify user about missed required metadata fields: id, updated_on
         # it is not required to have them, but it is recommended to have them for proper re-indexing and duplicate detection
@@ -292,7 +306,7 @@ class VectorStoreWrapper(BaseToolApiWrapper):
         logger.debug(f"Indexing documents: {documents}")
         logger.debug(self.vectoradapter)
 
-        documents = list(documents)
+        documents = documents + list(dependent_docs_generator)
         total_docs = len(documents)
         documents_count = 0
         _documents = []

alita_sdk/runtime/utils/streamlit.py

@@ -65,7 +65,8 @@ def pil_to_base64_string(pil_image):
         return None
 
 
-from alita_sdk.runtime.llms.alita import AlitaChatModel
+# from alita_sdk.runtime.llms.alita import AlitaChatModel
+from alita_sdk.runtime.clients.client import AlitaClient
 from alita_sdk.runtime.utils.AlitaCallback import AlitaStreamlitCallback
 from alita_sdk.runtime.toolkits.tools import get_toolkits, get_tools
 from alita_sdk.community.utils import check_schema
@@ -369,7 +370,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                 st.rerun()
         
         # Determine login form title and expansion state
-        if st.session_state.llm:
+        if st.session_state.client:
             login_title = "✅ Elitea Login (Connected)"
             # Collapse after successful login, but allow expansion
             if st.session_state.login_form_expanded is True:
@@ -385,13 +386,13 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
             deployment_secret = environ.get('XSECRET', 'secret')
             api_key_value = environ.get('API_KEY', None)
             project_id_value = int(environ.get('PROJECT_ID', 0))
-            if st.session_state.llm:
-                deployment_value = st.session_state.llm.deployment
-                api_key_value = st.session_state.llm.api_token
-                project_id_value = st.session_state.llm.project_id
-                
+            if st.session_state.client:
+                deployment_value = st.session_state.client.base_url
+                api_key_value = st.session_state.client.auth_token
+                project_id_value = st.session_state.client.project_id
+
             # Show current connection status
-            if st.session_state.llm:
+            if st.session_state.client:
                 st.success(f"Connected to: {deployment_value}")
                 st.info(f"Project ID: {project_id_value}")
                 
@@ -403,22 +404,24 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                 deployment_secret = st.text_input("Deployment Secret", placeholder="Enter Deployment Secret", value=deployment_secret)
                 
                 # Change button text based on login status
-                button_text = "Re-Login" if st.session_state.llm else "Login"
+                button_text = "Re-Login" if st.session_state.client else "Login"
                 submitted = st.form_submit_button(button_text)
                 
                 if submitted:
                     with st.spinner("Logging to Alita..."):
                         try:
-                            st.session_state.llm = AlitaChatModel(**{
-                                    "deployment": deployment,
-                                    "api_token": api_key,
-                                    "project_id": project_id,
-                                })
-                            client = st.session_state.llm.client
+
+                            st.session_state.client = AlitaClient(
+                                base_url=deployment,
+                                project_id=project_id,
+                                auth_token=api_key,
+                                api_extra_headers={"X-SECRET": deployment_secret}
+                            )
+                            
                             
                             # Fetch specific project secret for pgvector connection
                             try:
-                                pgvector_connstr = client.unsecret('pgvector_project_connstr')
+                                pgvector_connstr = st.session_state.client.unsecret('pgvector_project_connstr')
                                 if pgvector_connstr:
                                     st.session_state.project_secrets = {'pgvector_project_connstr': pgvector_connstr}
                                     logger.info("Successfully retrieved pgvector connection string from project secrets")
@@ -429,7 +432,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                 logger.warning(f"Could not retrieve pgvector connection string: {str(e)}")
                                 st.session_state.project_secrets = {}
                             
-                            integrations = client.all_models_and_integrations()
+                            integrations = st.session_state.client.all_models_and_integrations()
                             unique_models = set()
                             models_list = []
                             for entry in integrations:
@@ -438,7 +441,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                     if model.get('capabilities', {}).get('chat_completion') and model['name'] not in unique_models:
                                         unique_models.add(model['name'])
                                         models_list.append({'name': model['name'], 'integration_id': entry['uid']})
-                            st.session_state.agents = client.get_list_of_apps()
+                            st.session_state.agents = st.session_state.client.get_list_of_apps()
                             st.session_state.models = models_list
                             clear_chat_history()
                             
@@ -454,6 +457,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                         except Exception as e:
                             logger.error(f"Error loggin to ELITEA: {format_exc()}")
                             st.session_state.agents = None
+                            st.session_state.client = None
                             st.session_state.models = None
                             st.session_state.llm = None
                             st.session_state.project_secrets = None
@@ -463,7 +467,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
         llmconfig, toolkit_config = st.tabs(["Alita Agents", "Toolkit Testing"])
         
         with llmconfig:
-            if st.session_state.llm:
+            if st.session_state.client:
                 st.title("Available Agents")
                 st.write("This one will load latest version of agent")
                 with st.form("agents_form", clear_on_submit=False):
@@ -477,7 +481,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                             agent = next((a for a in st.session_state.agents if a['name'] == options), None)
                             if agent:
                                 agent_id = agent['id']
-                                agent_details = st.session_state.llm.client.get_app_details(agent_id)
+                                agent_details = st.session_state.client.get_app_details(agent_id)
                                 latest_version = next((v for v in agent_details['versions'] if v['name'] == agent_version_name), None)
                                 if latest_version:
                                     agent_version_id = latest_version['id']
@@ -504,11 +508,11 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                         
                                         # Try to get the complete agent configuration
                                         try:
-                                            agent_version_details = st.session_state.llm.client.get_app_version_details(agent_id, agent_version_id)
+                                            agent_version_details = st.session_state.client.get_app_version_details(agent_id, agent_version_id)
                                             agent_full_config = agent_version_details
                                         except AttributeError:
                                             try:
-                                                agent_version_details = st.session_state.llm.client.get_application_version_details(agent_id, agent_version_id)
+                                                agent_version_details = st.session_state.client.get_application_version_details(agent_id, agent_version_id)
                                                 agent_full_config = agent_version_details
                                             except AttributeError:
                                                 # Use the version details we already have
@@ -630,8 +634,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                                         st.session_state.agent_toolkit_configs = {}
                                         st.session_state.agent_raw_config = None
                                     
-                                    st.session_state.agent_executor = st.session_state.llm.client.application(
-                                        client=st.session_state.llm,
+                                    st.session_state.agent_executor = st.session_state.client.application(
                                         application_id=agent_id,
                                         application_version_id=agent_version_id,
                                         app_type=agent_type if agent_type else None,
@@ -658,7 +661,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
             """)
             
             # Check if user is logged in
-            if not st.session_state.llm:
+            if not st.session_state.client:
                 st.warning("⚠️ **Please log in first!**")
                 st.info("""
                 📋 **To use Toolkit Testing:**
@@ -671,7 +674,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                 st.stop()
             
             # User is logged in, proceed with toolkit testing
-            if st.session_state.llm:
+            if st.session_state.client:
                 # Show project secrets status with detailed debugging
                 secrets_status = st.session_state.project_secrets
                 
@@ -685,7 +688,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                 # Debug info (can be removed later)
                 with st.expander("🔍 Debug Info", expanded=False):
                     st.write(f"**Project Secrets Status:** {type(secrets_status)} - {secrets_status}")
-                    st.write(f"**LLM Status:** {'Connected' if st.session_state.llm else 'Not Connected'}")
+                    # st.write(f"**LLM Status:** {'Connected' if st.session_state.llm else 'Not Connected'}")
                 
                 # Toolkit selection and configuration
                 st.markdown("---")
@@ -1021,7 +1024,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                 st.markdown("👈 Please use the **Alita Login Form** in the sidebar to authenticate.")
                         
     # Main content area
-    if st.session_state.llm and st.session_state.agent_executor and st.session_state.agent_chat:
+    if st.session_state.client and st.session_state.agent_executor and st.session_state.agent_chat:
         try:
             st.title(st.session_state.agent_name)
         except:
@@ -1043,7 +1046,7 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
                 st.session_state.thread_id = response.get("thread_id", None)
                 st.session_state.messages.append({"role": "assistant", "content": response["output"]})
     
-    elif st.session_state.llm and st.session_state.show_toolkit_testing and st.session_state.configured_toolkit:
+    elif st.session_state.client and st.session_state.show_toolkit_testing and st.session_state.configured_toolkit:
         # Toolkit Testing Main View
         st.title("🚀 Toolkit Testing Interface")
         
@@ -1281,7 +1284,7 @@ Please explain how you would use these tools to help the user, even though I can
                 st.rerun()
     
     else:
-        if st.session_state.llm:
+        if st.session_state.client:
             st.title("🎯 Alita SDK Toolkit Interface")
             st.markdown("""
             ### Welcome to the Alita SDK!

alita_sdk/runtime/utils/utils.py

@@ -1,7 +1,12 @@
 import re
+from enum import Enum
 
 TOOLKIT_SPLITTER = "___"
 
+class IndexerKeywords(Enum):
+    DEPENDENT_DOCS = 'dependent_docs'
+    PARENT = 'parent_id'
+
 # This pattern matches characters that are NOT alphanumeric, underscores, or hyphens
 clean_string_pattern = re.compile(r'[^a-zA-Z0-9_.-]')
 

alita_sdk/tools/__init__.py

@@ -2,6 +2,7 @@ import logging
 from importlib import import_module
 from typing import Optional
 
+from langchain_core.tools import ToolException
 from langgraph.store.base import BaseStore
 
 logger = logging.getLogger(__name__)
@@ -113,6 +114,7 @@ def get_tools(tools_list, alita, llm, store: Optional[BaseStore] = None, *args,
 
             except Exception as e:
                 logger.error(f"Error getting tools for {tool_type}: {e}")
+                raise ToolException(f"Error getting tools for {tool_type}: {e}")
 
         # Handle ADO repos special case (it might be requested as azure_devops_repos)
         elif tool_type in ['ado_repos', 'azure_devops_repos'] and 'ado_repos' in AVAILABLE_TOOLS:

alita_sdk/tools/ado/repos/repos_wrapper.py

@@ -250,6 +250,7 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
     token: Optional[SecretStr]
     _client: Optional[GitClient] = PrivateAttr()
 
+    llm: Optional[Any] = None
     # Vector store configuration
     connection_string: Optional[SecretStr] = None
     collection_name: Optional[str] = None
@@ -303,24 +304,30 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
 
     def _get_files(
             self,
-            directory_path: str = "",
-            branch_name: str = None,
+            path: str = "",
+            branch: str = None,
             recursion_level: str = "Full",
     ) -> str:
+        """Get list of files from a repository path and branch.
+
+        Args:
+            path (str): Path within the repository to list files from
+            branch (str): Branch to get files from. Defaults to base_branch if None.
+            recursion_level (str): OneLevel - includes immediate children, Full - includes all items, None - no recursion
+
+        Returns:
+            List[str]: List of file paths
         """
-        Params:
-            recursion_level: OneLevel - includes immediate children, Full - includes all items, None - no recursion
-        """
-        branch_name = branch_name if branch_name else self.base_branch
+        branch = branch if branch else self.base_branch
         files: List[str] = []
         try:
             version_descriptor = GitVersionDescriptor(
-                version=branch_name, version_type="branch"
+                version=branch, version_type="branch"
             )
             items = self._client.get_items(
                 repository_id=self.repository_id,
                 project=self.project,
-                scope_path=directory_path,
+                scope_path=path,
                 recursion_level=recursion_level,
                 version_descriptor=version_descriptor,
                 include_content_metadata=True,
@@ -334,7 +341,7 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
             item = items.pop(0)
             if item.git_object_type == "blob":
                 files.append(item.path)
-        return str(files)
+        return files # Changed to return list directly instead of str
 
     def set_active_branch(self, branch_name: str) -> str:
         """
@@ -389,7 +396,7 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
             logger.error(msg)
             return ToolException(msg)
 
-    def list_files(self, directory_path: str = "", branch_name: str = None) -> str:
+    def list_files(self, directory_path: str = "", branch_name: str = None) -> List[str]:
         """
         Recursively fetches files from a directory in the repo.
 
@@ -398,12 +405,12 @@ class ReposApiWrapper(BaseCodeToolApiWrapper):
             branch_name (str): The name of the branch where the files to be received.
 
         Returns:
-            str: List of file paths, or an error message.
+            List[str]: List of file paths, or an error message.
         """
         self.active_branch = branch_name if branch_name else self.active_branch
         return self._get_files(
-            directory_path=directory_path,
-            branch_name=self.active_branch if self.active_branch else self.base_branch,
+            path=directory_path,
+            branch=self.active_branch if self.active_branch else self.base_branch,
         )
 
     def parse_pull_request_comments(