import os import base64 import json import requests import asyncio from enum import Enum from typing import List, Dict, Optional, Callable, Awaitable from pydantic import BaseModel, Field, ValidationError from sentence_transformers import SentenceTransformer from markdownify import markdownify import numpy as np from rank_bm25 import BM25Okapi from sklearn.neighbors import NearestNeighbors # ----------------------------------------------------------------------------- # Configuration models # ----------------------------------------------------------------------------- class Valves(BaseModel): base_url: str = Field("https://example.atlassian.net/wiki", description="Base URL of Confluence") username: str = Field("", description="Username/email for API key auth") api_key: str = Field("", description="API key or personal access token") api_key_auth: bool = Field(True, description="Use API key auth; if False, treat api_key as PAT") ssl_verify: bool = Field(True, description="Verify SSL certificates") api_result_limit: int = Field(5, ge=1, description="Max pages to fetch from API") class UserValves(BaseModel): api_key_auth: Optional[bool] username: Optional[str] api_key: Optional[str] split_terms: Optional[bool] included_confluence_spaces: Optional[str] excluded_confluence_spaces: Optional[str] # ----------------------------------------------------------------------------- # Helper functions # ----------------------------------------------------------------------------- def parse_user_valves(user_dict: Dict, global_valves: Valves) -> Dict: """ Normalize user overrides, merging with global default valves. """ try: # Accept either flat or nested under 'valves' raw = user_dict.get('valves') if user_dict.get('valves') else user_dict uv = UserValves(**raw) except (ValidationError, TypeError): return { 'api_key_auth': global_valves.api_key_auth, 'username': global_valves.username, 'api_key': global_valves.api_key, 'split_terms': True, 'included_confluence_spaces': None, 'excluded_confluence_spaces': None, } return { 'api_key_auth': uv.api_key_auth if uv.api_key_auth is not None else global_valves.api_key_auth, 'username': uv.username or global_valves.username, 'api_key': uv.api_key or global_valves.api_key, 'split_terms': uv.split_terms if uv.split_terms is not None else True, 'included_confluence_spaces': uv.included_confluence_spaces, 'excluded_confluence_spaces': uv.excluded_confluence_spaces, } # ----------------------------------------------------------------------------- # Confluence client # ----------------------------------------------------------------------------- class ConfluenceError(Exception): pass class ConfluenceAuthError(ConfluenceError): pass class ConfluenceAPIError(ConfluenceError): pass class ConfluenceClient: def __init__(self, base_url: str, username: str, api_key: str, api_key_auth: bool, ssl_verify: bool): self.base_url = base_url.rstrip('/') self.headers = self._make_headers(username, api_key, api_key_auth) self.ssl_verify = ssl_verify def _make_headers(self, username: str, api_key: str, api_key_auth: bool) -> Dict[str,str]: if api_key_auth: auth = base64.b64encode(f"{username}:{api_key}".encode()).decode() return {'Authorization': f'Basic {auth}'} return {'Authorization': f'Bearer {api_key}'} def get(self, endpoint: str, params: Dict = None) -> Dict: url = f"{self.base_url}/rest/api/{endpoint}" try: resp = requests.get(url, headers=self.headers, params=params, verify=self.ssl_verify) except requests.RequestException as e: raise ConfluenceAPIError(str(e)) if resp.status_code == 401: raise ConfluenceAuthError("Authentication failed.") if not resp.ok: raise ConfluenceAPIError(resp.text) return resp.json() def search(self, cql: str, limit: int) -> List[Dict]: data = self.get('content/search', {'cql': cql, 'limit': limit}) return data.get('results', []) def get_page(self, page_id: str) -> Dict: data = self.get(f'content/{page_id}', {'expand': 'body.view'}) body_html = data['body']['view']['value'] body_md = markdownify(body_html) return { 'id': data['id'], 'title': data['title'], 'body': body_md, 'link': f"{self.base_url}{data['_links']['webui']}" } # ----------------------------------------------------------------------------- # Main Tool # ----------------------------------------------------------------------------- class Tools: def __init__(self): self.valves = Valves() async def search_confluence( self, query: str, type: str, event_emitter: Callable[[dict], Awaitable[None]], __user__: Dict = {} ) -> str: """ Search Confluence pages by title, content, or both. Returns JSON list of {{id, title, body, link}}. """ # Normalize user settings settings = parse_user_valves(__user__, self.valves) # Validate credentials if settings['api_key_auth'] and not settings['username'] or not settings['api_key']: await event_emitter({'type':'status','data':{'description':'❌ Missing username or API key','status':'complete','done':True}}) return json.dumps([]) # Build CQL field = None if type == 'title': field='title' elif type == 'content': field='text' terms = query.split() if settings['split_terms'] else [query] clauses = [f"{field or '(title OR text)'} ~ \"{term}\"" for term in terms] spaces_inc = settings['included_confluence_spaces'] spaces_exc = settings['excluded_confluence_spaces'] if spaces_inc: clauses.append(f"space in ({','.join(f"'{s.strip()}'" for s in spaces_inc.split(','))})") if spaces_exc: clauses.append(f"space not in ({','.join(f"'{s.strip()}'" for s in spaces_exc.split(','))})") cql = ' AND '.join(clauses) # Fetch list of page IDs try: pages = ConfluenceClient( self.valves.base_url, settings['username'], settings['api_key'], settings['api_key_auth'], self.valves.ssl_verify ).search(cql, self.valves.api_result_limit) except ConfluenceError as e: await event_emitter({'type':'status','data':{'description':f'❌ {e}','status':'complete','done':True}}) return json.dumps([]) # Retrieve each page results = [] client = ConfluenceClient( self.valves.base_url, settings['username'], settings['api_key'], settings['api_key_auth'], self.valves.ssl_verify ) for idx, page in enumerate(pages, 1): await event_emitter({'type':'status','data':{'description':f'🔎 Retrieving page {idx}/{len(pages)}','status':'in_progress','done':False}}) pg = client.get_page(page['id']) results.append(pg) await event_emitter({'type':'status','data':{'description':f'✅ Found {len(results)} pages','status':'complete','done':True}}) return json.dumps(results)