Add support for documentation content exceeding token limits (#1670)

* - Add support for documentation content exceeding token limits via two phase operation:
1. Ask LLM to rank headings which are most likely to contain an answer to a user question
2. Provide the corresponding files for the LLM to search for an answer.

- Refactor of help_docs to make the code more readable
- For the purpose of getting canonical path: git providers to use default branch and not the PR's source branch.
- Refactor of token counting and making it clear on when an estimate factor will be used.

* Code review changes:
1. Correctly handle exception during retry_with_fallback_models (to allow fallback model to run in case of failure)
2. Better naming for default_branch in bitbucket cloud provider

pr_agent/git_providers/bitbucket_provider.py

@@ -92,7 +92,7 @@ class BitbucketProvider(GitProvider):
                 return ("", "")
             workspace_name, project_name = repo_path.split('/')
         else:
-            desired_branch = self.get_pr_branch()
+            desired_branch = self.get_repo_default_branch()
             parsed_pr_url = urlparse(self.pr_url)
             scheme_and_netloc = parsed_pr_url.scheme + "://" + parsed_pr_url.netloc
             workspace_name, project_name = (self.workspace_slug, self.repo_slug)
@@ -470,6 +470,16 @@ class BitbucketProvider(GitProvider):
     def get_pr_branch(self):
         return self.pr.source_branch
 
+    # This function attempts to get the default branch of the repository. As a fallback, uses the PR destination branch.
+    # Note: Must be running from a PR context.
+    def get_repo_default_branch(self):
+        try:
+            url_repo = f"https://api.bitbucket.org/2.0/repositories/{self.workspace_slug}/{self.repo_slug}/"
+            response_repo = requests.request("GET", url_repo, headers=self.headers).json()
+            return response_repo['mainbranch']['name']
+        except:
+            return self.pr.destination_branch
+
     def get_pr_owner_id(self) -> str | None:
         return self.workspace_slug
 

pr_agent/git_providers/bitbucket_server_provider.py

@@ -64,9 +64,15 @@ class BitbucketServerProvider(GitProvider):
         workspace_name = None
         project_name = None
         if not repo_git_url:
-            desired_branch = self.get_pr_branch()
             workspace_name = self.workspace_slug
             project_name = self.repo_slug
+            default_branch_dict = self.bitbucket_client.get_default_branch(workspace_name, project_name)
+            if 'displayId' in default_branch_dict:
+                desired_branch = default_branch_dict['displayId']
+            else:
+                get_logger().error(f"Cannot obtain default branch for workspace_name={workspace_name}, "
+                                   f"project_name={project_name}, default_branch_dict={default_branch_dict}")
+                return ("", "")
         elif '.git' in repo_git_url and 'scm/' in repo_git_url:
             repo_path = repo_git_url.split('.git')[0].split('scm/')[-1]
             if repo_path.count('/') == 1:  # Has to have the form <workspace>/<repo>

pr_agent/git_providers/github_provider.py

@@ -133,7 +133,7 @@ class GithubProvider(GitProvider):
         if (not owner or not repo) and self.repo: #"else" - User did not provide an external git url, or not an issue, use self.repo object
             owner, repo = self.repo.split('/')
             scheme_and_netloc = self.base_url_html
-            desired_branch = self.get_pr_branch()
+            desired_branch = self.repo_obj.default_branch
         if not all([scheme_and_netloc, owner, repo]): #"else": Not invoked from a PR context,but no provided git url for context
             get_logger().error(f"Unable to get canonical url parts since missing context (PR or explicit git url)")
             return ("", "")

pr_agent/git_providers/gitlab_provider.py

@@ -87,7 +87,11 @@ class GitLabProvider(GitProvider):
             return ("", "")
         if not repo_git_url: #Use PR url as context
             repo_path = self._get_project_path_from_pr_or_issue_url(self.pr_url)
-            desired_branch = self.get_pr_branch()
+            try:
+                desired_branch = self.gl.projects.get(self.id_project).default_branch
+            except Exception as e:
+                get_logger().exception(f"Cannot get PR: {self.pr_url} default branch. Tried project ID: {self.id_project}")
+                return ("", "")
         else: #Use repo git url
             repo_path = repo_git_url.split('.git')[0].split('.com/')[-1]
         prefix = f"{self.gitlab_url}/{repo_path}/-/blob/{desired_branch}"