adaptive-memory-engine 0.1.6__py3-none-any.whl

ame/connectors/json_helpers.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+
+def read_json(path: Path) -> Any:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def as_list(value: Any) -> list[Any]:
+    if isinstance(value, list):
+        return value
+    if isinstance(value, dict):
+        for key in [
+            "messages",
+            "issues",
+            "pull_requests",
+            "prs",
+            "discussions",
+            "comments",
+            "files",
+            "documents",
+            "threads",
+            "events",
+            "rows",
+            "values",
+            "items",
+        ]:
+            rows = value.get(key)
+            if isinstance(rows, list):
+                return rows
+        return [value]
+    return []
+
+
+def first_present(row: dict[str, Any], *keys: str) -> Any:
+    for key in keys:
+        value = row.get(key)
+        if value is not None:
+            return value
+    return None

ame/connectors/markdown.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import hashlib
+import re
+from pathlib import Path
+from typing import Any
+
+from ame.bronze.schema import BronzeDocument
+from ame.connectors.base import SourceRef
+
+
+FRONTMATTER_RE = re.compile(r"\A---\s*\n(.*?)\n---\s*\n", re.DOTALL)
+HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)
+
+
+class MarkdownConnector:
+    source_type = "markdown"
+
+    def scan(self, path: Path) -> list[SourceRef]:
+        root = path.expanduser().resolve()
+        files = [root] if root.is_file() else sorted([*root.rglob("*.md"), *root.rglob("*.markdown")])
+        refs: list[SourceRef] = []
+        for file in files:
+            source_id = str(file.relative_to(root) if root.is_dir() else file.name)
+            content = file.read_text(encoding="utf-8")
+            sections = self._sections(content)
+            if len(sections) <= 1:
+                refs.append(SourceRef(path=file, source_id=source_id, content=content, root_source_id=source_id))
+                continue
+            for index, section in enumerate(sections):
+                refs.append(
+                    SourceRef(
+                        path=file,
+                        source_id=f"{source_id}#{self._safe_section_id(section['title'])}",
+                        content=section["content"],
+                        section_title=section["title"],
+                        section_path=tuple(section["path"]),
+                        section_level=section["level"],
+                        section_index=index,
+                        root_source_id=source_id,
+                    )
+                )
+        return refs
+
+    def load(self, corpus_id: str, ref: SourceRef) -> BronzeDocument:
+        content = ref.content if ref.content is not None else ref.path.read_text(encoding="utf-8")
+        digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
+        metadata = self._metadata(content, ref.path)
+        if ref.section_title:
+            metadata["title"] = ref.section_title
+            metadata["section_title"] = ref.section_title
+            metadata["section_path"] = list(ref.section_path)
+            metadata["section_level"] = ref.section_level
+            metadata["section_index"] = ref.section_index
+            metadata["root_source_id"] = ref.root_source_id or ref.source_id.split("#", 1)[0]
+            metadata["source_file"] = ref.root_source_id or ref.source_id.split("#", 1)[0]
+        return BronzeDocument(
+            id=f"bronze_{digest[:16]}",
+            corpus_id=corpus_id,
+            source_type=self.source_type,
+            source_id=ref.source_id,
+            content=content,
+            metadata=metadata,
+            content_hash=f"sha256:{digest}",
+        )
+
+    def _metadata(self, content: str, path: Path) -> dict[str, Any]:
+        metadata: dict[str, Any] = {"path": str(path), "title": path.stem}
+        frontmatter = FRONTMATTER_RE.search(content)
+        if frontmatter:
+            metadata["frontmatter"] = self._parse_frontmatter(frontmatter.group(1))
+            if title := metadata["frontmatter"].get("title"):
+                metadata["title"] = title
+        metadata["headings"] = [match.group(2).strip() for match in HEADING_RE.finditer(content)]
+        if metadata["title"] == path.stem and metadata["headings"]:
+            metadata["title"] = metadata["headings"][0]
+        return metadata
+
+    def _parse_frontmatter(self, text: str) -> dict[str, str]:
+        data: dict[str, str] = {}
+        for line in text.splitlines():
+            if ":" not in line:
+                continue
+            key, value = line.split(":", 1)
+            data[key.strip()] = value.strip().strip('"')
+        return data
+
+    def _sections(self, content: str) -> list[dict]:
+        matches = [match for match in HEADING_RE.finditer(content) if len(match.group(1)) <= 2]
+        if len(matches) <= 1:
+            return []
+        sections: list[dict] = []
+        path_by_level: dict[int, str] = {}
+        for index, match in enumerate(matches):
+            level = len(match.group(1))
+            title = match.group(2).strip()
+            path_by_level[level] = title
+            for deeper in [known for known in path_by_level if known > level]:
+                del path_by_level[deeper]
+            start = match.start()
+            end = matches[index + 1].start() if index + 1 < len(matches) else len(content)
+            section_content = content[start:end].strip()
+            if not section_content:
+                continue
+            sections.append(
+                {
+                    "title": title,
+                    "level": level,
+                    "path": [path_by_level[key] for key in sorted(path_by_level)],
+                    "content": section_content,
+                }
+            )
+        return sections
+
+    def _safe_section_id(self, title: str) -> str:
+        return re.sub(r"\s+", " ", title).strip().replace("/", "-").replace(":", " -")

ame/connectors/notion.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+from typing import Any
+
+from ame.bronze.schema import BronzeDocument
+from ame.connectors.base import SourceRef
+from ame.connectors.json_helpers import as_list, first_present, read_json
+
+
+class NotionConnector:
+    source_type = "notion"
+
+    def scan(self, path: Path) -> list[SourceRef]:
+        root = path.expanduser().resolve()
+        files = [root] if root.is_file() else sorted(root.rglob("*.json"))
+        refs: list[SourceRef] = []
+        for file in files:
+            for row in as_list(read_json(file)):
+                if not isinstance(row, dict):
+                    continue
+                page_id = str(first_present(row, "id", "page_id") or file.stem)
+                refs.append(SourceRef(path=file, source_id=f"notion:{page_id}", content=self._page_content(row, page_id)))
+        return refs
+
+    def load(self, corpus_id: str, ref: SourceRef) -> BronzeDocument:
+        content = ref.content or ref.path.read_text(encoding="utf-8")
+        digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
+        return BronzeDocument(
+            id=f"bronze_{digest[:16]}",
+            corpus_id=corpus_id,
+            source_type=self.source_type,
+            source_id=ref.source_id,
+            content=content,
+            metadata={"path": str(ref.path), "title": ref.source_id, "connector": "notion"},
+            content_hash=f"sha256:{digest}",
+        )
+
+    def _page_content(self, row: dict[str, Any], page_id: str) -> str:
+        title = str(first_present(row, "title", "name") or page_id)
+        body = str(first_present(row, "content", "body", "text") or "")
+        url = str(first_present(row, "url", "public_url") or "")
+        last_edited_time = str(first_present(row, "last_edited_time", "updated_at") or "")
+        return "\n".join(
+            [
+                "---",
+                f"title: {title}",
+                f"notion_page_id: {page_id}",
+                f"url: {url}",
+                f"last_edited_time: {last_edited_time}",
+                "---",
+                "",
+                f"# {title}",
+                "",
+                body,
+                "",
+            ]
+        )

ame/connectors/oauth_callback.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import secrets
+import urllib.parse
+import webbrowser
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+
+class OAuthCallbackError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class OAuthCallbackResult:
+    code: str
+    state: str
+
+
+def new_oauth_state() -> str:
+    return secrets.token_urlsafe(24)
+
+
+def run_local_oauth_login(
+    authorization_url: str,
+    redirect_uri: str,
+    expected_state: str,
+    *,
+    open_browser: bool = True,
+    timeout_seconds: int = 180,
+) -> OAuthCallbackResult:
+    parsed = urllib.parse.urlparse(redirect_uri)
+    if parsed.scheme != "http" or parsed.hostname not in {"localhost", "127.0.0.1"}:
+        raise OAuthCallbackError("Local OAuth login requires an http://localhost redirect URI")
+    if not parsed.port:
+        raise OAuthCallbackError("Local OAuth login redirect URI must include a port")
+
+    server = _OAuthCallbackServer((parsed.hostname, parsed.port), _OAuthCallbackHandler)
+    server.expected_state = expected_state
+    server.timeout = timeout_seconds
+    try:
+        if open_browser:
+            webbrowser.open(authorization_url)
+        server.handle_request()
+    finally:
+        server.server_close()
+
+    if server.error:
+        raise OAuthCallbackError(server.error)
+    if server.result is None:
+        raise OAuthCallbackError("OAuth login timed out before receiving a callback")
+    return server.result
+
+
+class _OAuthCallbackServer(HTTPServer):
+    expected_state: str
+    result: OAuthCallbackResult | None = None
+    error: str | None = None
+
+
+class _OAuthCallbackHandler(BaseHTTPRequestHandler):
+    def do_GET(self) -> None:
+        parsed = urllib.parse.urlparse(self.path)
+        query = urllib.parse.parse_qs(parsed.query)
+        state = _single(query, "state")
+        code = _single(query, "code")
+        error = _single(query, "error")
+
+        if error:
+            self.server.error = f"OAuth provider returned error: {error}"  # type: ignore[attr-defined]
+            self._respond(400, "OAuth login failed. You can close this window.")
+            return
+        if state != self.server.expected_state:  # type: ignore[attr-defined]
+            self.server.error = "OAuth callback state did not match"  # type: ignore[attr-defined]
+            self._respond(400, "OAuth login state mismatch. You can close this window.")
+            return
+        if not code:
+            self.server.error = "OAuth callback did not include code"  # type: ignore[attr-defined]
+            self._respond(400, "OAuth login did not include a code. You can close this window.")
+            return
+
+        self.server.result = OAuthCallbackResult(code=code, state=state)  # type: ignore[attr-defined]
+        self._respond(200, "OAuth login complete. You can close this window.")
+
+    def log_message(self, format: str, *args) -> None:
+        return
+
+    def _respond(self, status: int, message: str) -> None:
+        payload = (
+            "<!doctype html><html><head><meta charset=\"utf-8\"><title>Adaptive Memory Engine</title></head>"
+            f"<body><h1>{message}</h1></body></html>"
+        ).encode("utf-8")
+        self.send_response(status)
+        self.send_header("Content-Type", "text/html; charset=utf-8")
+        self.send_header("Content-Length", str(len(payload)))
+        self.end_headers()
+        self.wfile.write(payload)
+
+
+def _single(query: dict[str, list[str]], key: str) -> str:
+    values = query.get(key) or []
+    return values[0] if values else ""

ame/connectors/oauth_provider.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+import base64
+import json
+import urllib.parse
+import urllib.request
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Literal, Protocol
+
+from pydantic import BaseModel, Field
+
+from ame.core.paths import ame_home
+from ame.security import token_vault
+
+
+TokenRequestStyle = Literal["form", "json"]
+ClientAuthStyle = Literal["body", "basic"]
+
+
+class ConnectedAppOAuthError(RuntimeError):
+    pass
+
+
+class OAuthProviderSpec(BaseModel):
+    name: str
+    authorize_url: str
+    token_url: str
+    default_redirect_uri: str
+    default_scopes: list[str] = Field(default_factory=list)
+    scope_separator: str = " "
+    auth_extra: dict[str, str] = Field(default_factory=dict)
+    token_request_style: TokenRequestStyle = "form"
+    client_auth_style: ClientAuthStyle = "body"
+
+
+class ConnectedAppOAuthConfig(BaseModel):
+    provider: str
+    client_id: str = ""
+    client_secret: str = ""
+    redirect_uri: str
+    scopes: list[str] = Field(default_factory=list)
+
+
+class ConnectedAppToken(BaseModel):
+    provider: str
+    account_id: str
+    access_token: str
+    refresh_token: str | None = None
+    token_type: str = "Bearer"
+    expires_in: int | None = None
+    scopes: list[str] = Field(default_factory=list)
+    raw: dict[str, Any] = Field(default_factory=dict)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+
+
+class ConnectedAppHttpClient(Protocol):
+    def post_json(
+        self,
+        url: str,
+        data: dict[str, Any],
+        headers: dict[str, str] | None = None,
+        *,
+        json_body: bool = False,
+    ) -> dict[str, Any]:
+        ...
+
+
+class UrlLibConnectedAppHttpClient:
+    def post_json(
+        self,
+        url: str,
+        data: dict[str, Any],
+        headers: dict[str, str] | None = None,
+        *,
+        json_body: bool = False,
+    ) -> dict[str, Any]:
+        request_headers = dict(headers or {})
+        if json_body:
+            body = json.dumps(data).encode("utf-8")
+            request_headers.setdefault("Content-Type", "application/json")
+        else:
+            body = urllib.parse.urlencode(data).encode("utf-8")
+            request_headers.setdefault("Content-Type", "application/x-www-form-urlencoded")
+        request_headers.setdefault("Accept", "application/json")
+        request = urllib.request.Request(url, data=body, headers=request_headers, method="POST")
+        with urllib.request.urlopen(request, timeout=30) as response:
+            return json.loads(response.read().decode("utf-8"))
+
+
+PROVIDER_SPECS: dict[str, OAuthProviderSpec] = {
+    "github": OAuthProviderSpec(
+        name="github",
+        authorize_url="https://github.com/login/oauth/authorize",
+        token_url="https://github.com/login/oauth/access_token",
+        default_redirect_uri="http://localhost:8765/github/oauth/callback",
+        default_scopes=["read:user"],
+    ),
+    "notion": OAuthProviderSpec(
+        name="notion",
+        authorize_url="https://api.notion.com/v1/oauth/authorize",
+        token_url="https://api.notion.com/v1/oauth/token",
+        default_redirect_uri="http://localhost:8765/notion/oauth/callback",
+        default_scopes=[],
+        auth_extra={"owner": "user"},
+        token_request_style="json",
+        client_auth_style="basic",
+    ),
+    "jira": OAuthProviderSpec(
+        name="jira",
+        authorize_url="https://auth.atlassian.com/authorize",
+        token_url="https://auth.atlassian.com/oauth/token",
+        default_redirect_uri="http://localhost:8765/jira/oauth/callback",
+        default_scopes=["read:jira-work", "read:jira-user", "offline_access"],
+        auth_extra={"audience": "api.atlassian.com", "prompt": "consent"},
+        token_request_style="json",
+    ),
+}
+
+
+def provider_spec(provider: str) -> OAuthProviderSpec:
+    normalized = provider.casefold().replace("_", "-")
+    aliases = {"atlassian": "jira", "github-oauth": "github", "notion-oauth": "notion", "jira-oauth": "jira"}
+    normalized = aliases.get(normalized, normalized)
+    spec = PROVIDER_SPECS.get(normalized)
+    if not spec:
+        raise ConnectedAppOAuthError(f"Unsupported connected app provider: {provider}")
+    return spec
+
+
+class ConnectedAppOAuthClient:
+    def __init__(self, config: ConnectedAppOAuthConfig, spec: OAuthProviderSpec | None = None, http: ConnectedAppHttpClient | None = None):
+        self.config = config
+        self.spec = spec or provider_spec(config.provider)
+        self.http = http or UrlLibConnectedAppHttpClient()
+
+    def authorization_url(self, state: str) -> str:
+        params = {
+            "client_id": self.config.client_id,
+            "redirect_uri": self.config.redirect_uri,
+            "response_type": "code",
+            "state": state,
+            **self.spec.auth_extra,
+        }
+        if self.config.scopes:
+            params["scope"] = self.spec.scope_separator.join(self.config.scopes)
+        return f"{self.spec.authorize_url}?{urllib.parse.urlencode(params)}"
+
+    def exchange_code(self, code: str, account_id: str = "default") -> ConnectedAppToken:
+        headers: dict[str, str] = {}
+        payload = {
+            "code": code,
+            "redirect_uri": self.config.redirect_uri,
+            "grant_type": "authorization_code",
+        }
+        if self.spec.client_auth_style == "basic":
+            auth = f"{self.config.client_id}:{self.config.client_secret}".encode("utf-8")
+            headers["Authorization"] = f"Basic {base64.b64encode(auth).decode('ascii')}"
+        else:
+            payload["client_id"] = self.config.client_id
+            payload["client_secret"] = self.config.client_secret
+
+        response = self.http.post_json(
+            self.spec.token_url,
+            payload,
+            headers=headers,
+            json_body=self.spec.token_request_style == "json",
+        )
+        if error := response.get("error"):
+            raise ConnectedAppOAuthError(f"{self.spec.name} OAuth exchange failed: {error}")
+        access_token = str(response.get("access_token") or "")
+        if not access_token:
+            raise ConnectedAppOAuthError(f"{self.spec.name} OAuth response did not include access_token")
+        return ConnectedAppToken(
+            provider=self.spec.name,
+            account_id=account_id,
+            access_token=access_token,
+            refresh_token=response.get("refresh_token"),
+            token_type=str(response.get("token_type") or "Bearer"),
+            expires_in=response.get("expires_in"),
+            scopes=_split_scopes(response.get("scope")) or list(self.config.scopes),
+            raw={key: value for key, value in response.items() if key not in {"access_token", "refresh_token"}},
+        )
+
+
+class ConnectedAppTokenStore:
+    def __init__(self, provider: str, path: Path | None = None, backend: str = "file"):
+        self.provider = provider_spec(provider).name
+        self.path = path or ame_home() / "tokens" / f"{self.provider}.json"
+        self.vault = token_vault(self.provider, self.path, backend=backend)  # type: ignore[arg-type]
+
+    def save(self, token: ConnectedAppToken) -> ConnectedAppToken:
+        data = self._read()
+        data[token.account_id] = token.model_dump(mode="json")
+        self.vault.save(data)
+        return token
+
+    def load(self, account_id: str = "default") -> ConnectedAppToken:
+        data = self._read()
+        row = data.get(account_id)
+        if not isinstance(row, dict):
+            raise ConnectedAppOAuthError(f"{self.provider} token not found for account_id={account_id}")
+        return ConnectedAppToken.model_validate(row)
+
+    def revoke(self, account_id: str = "default") -> bool:
+        data = self._read()
+        existed = account_id in data
+        data.pop(account_id, None)
+        if data:
+            self.vault.save(data)
+        else:
+            self.vault.delete()
+        return existed
+
+    def _read(self) -> dict[str, Any]:
+        return self.vault.load()
+
+
+def exchange_and_save_connected_app_token(
+    provider: str,
+    code: str,
+    config: ConnectedAppOAuthConfig,
+    *,
+    account_id: str = "default",
+    store_path: Path | None = None,
+    token_backend: str = "file",
+    http: ConnectedAppHttpClient | None = None,
+) -> ConnectedAppToken:
+    spec = provider_spec(provider)
+    token = ConnectedAppOAuthClient(config, spec=spec, http=http).exchange_code(code, account_id=account_id)
+    return ConnectedAppTokenStore(spec.name, store_path, backend=token_backend).save(token)
+
+
+def default_config(provider: str, *, client_id: str = "", client_secret: str = "", redirect_uri: str | None = None, scopes: list[str] | None = None) -> ConnectedAppOAuthConfig:
+    spec = provider_spec(provider)
+    return ConnectedAppOAuthConfig(
+        provider=spec.name,
+        client_id=client_id,
+        client_secret=client_secret,
+        redirect_uri=redirect_uri or spec.default_redirect_uri,
+        scopes=list(spec.default_scopes if scopes is None else scopes),
+    )
+
+
+def _split_scopes(value: Any) -> list[str]:
+    if not value:
+        return []
+    if isinstance(value, list):
+        return [str(item) for item in value if str(item).strip()]
+    return [item for item in str(value).replace(",", " ").split() if item]

ame/connectors/obsidian.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import re
+
+from ame.connectors.markdown import MarkdownConnector
+
+
+WIKILINK_RE = re.compile(r"\[\[([^\]]+)\]\]")
+TAG_RE = re.compile(r"(?<!\w)#([A-Za-z0-9_/-]+)")
+
+
+class ObsidianConnector(MarkdownConnector):
+    source_type = "obsidian"
+
+    def _metadata(self, content: str, path):
+        metadata = super()._metadata(content, path)
+        metadata["wikilinks"] = sorted({m.group(1).split("|", 1)[0].strip() for m in WIKILINK_RE.finditer(content)})
+        metadata["tags"] = sorted({m.group(1) for m in TAG_RE.finditer(content)})
+        return metadata