From 713aac97e14fe5fcbc3a3072b55f0e7aa10f1e84 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Tue, 28 Jul 2026 10:00:37 +0800 Subject: [PATCH 1/4] refactor: move all pydantic models to schemas module, and add async version of save and load --- api/api.py | 156 ++++--------------------- api/schemas/__init__.py | 36 ++++++ api/{chat_model.py => schemas/chat.py} | 40 ++++--- api/schemas/io.py | 20 ++++ api/schemas/models.py | 31 +++++ api/schemas/wiki.py | 108 +++++++++++++++++ api/simple_chat.py | 2 +- api/websocket_wiki.py | 2 +- 8 files changed, 248 insertions(+), 147 deletions(-) create mode 100644 api/schemas/__init__.py rename api/{chat_model.py => schemas/chat.py} (60%) create mode 100644 api/schemas/io.py create mode 100644 api/schemas/models.py create mode 100644 api/schemas/wiki.py diff --git a/api/api.py b/api/api.py index d40e73f96..19be43250 100644 --- a/api/api.py +++ b/api/api.py @@ -1,18 +1,31 @@ -import os +import asyncio +import json import logging -from fastapi import FastAPI, HTTPException, Query, Request, WebSocket +import os +from datetime import datetime +from typing import List, Optional + +from fastapi import FastAPI, HTTPException, Query from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, Response -from typing import List, Optional, Dict, Any, Literal -import json -from datetime import datetime -from pydantic import BaseModel, Field -import google.generativeai as genai -import asyncio -# Configure logging +from api.config import WIKI_AUTH_CODE, WIKI_AUTH_MODE, configs from api.logging_config import setup_logging +from api.schemas import ( + AuthorizationConfig, + Model, + ModelConfig, + ProcessedProjectEntry, + Provider, + WikiCacheData, + WikiCacheRequest, + WikiExportRequest, + WikiPage, + aload, + asave, +) +# Configure logging setup_logging() logger = logging.getLogger(__name__) @@ -36,115 +49,6 @@ def get_adalflow_default_root_path(): return os.path.expanduser(os.path.join("~", ".adalflow")) -# --- Pydantic Models --- -class WikiPage(BaseModel): - """ - Model for a wiki page. - """ - id: str - title: str - content: str - filePaths: List[str] - importance: str # Should ideally be Literal['high', 'medium', 'low'] - relatedPages: List[str] - -class ProcessedProjectEntry(BaseModel): - id: str # Filename - owner: str - repo: str - name: str # owner/repo - repo_type: str # Renamed from type to repo_type for clarity with existing models - submittedAt: int # Timestamp - language: str # Extracted from filename - -class RepoInfo(BaseModel): - owner: str - repo: str - type: str - token: Optional[str] = None - localPath: Optional[str] = None - repoUrl: Optional[str] = None - - -class WikiSection(BaseModel): - """ - Model for the wiki sections. - """ - id: str - title: str - pages: List[str] - subsections: Optional[List[str]] = None - - -class WikiStructureModel(BaseModel): - """ - Model for the overall wiki structure. - """ - id: str - title: str - description: str - pages: List[WikiPage] - sections: Optional[List[WikiSection]] = None - rootSections: Optional[List[str]] = None - -class WikiCacheData(BaseModel): - """ - Model for the data to be stored in the wiki cache. - """ - wiki_structure: WikiStructureModel - generated_pages: Dict[str, WikiPage] - repo_url: Optional[str] = None #compatible for old cache - repo: Optional[RepoInfo] = None - provider: Optional[str] = None - model: Optional[str] = None - -class WikiCacheRequest(BaseModel): - """ - Model for the request body when saving wiki cache. - """ - repo: RepoInfo - language: str - wiki_structure: WikiStructureModel - generated_pages: Dict[str, WikiPage] - provider: str - model: str - -class WikiExportRequest(BaseModel): - """ - Model for requesting a wiki export. - """ - repo_url: str = Field(..., description="URL of the repository") - pages: List[WikiPage] = Field(..., description="List of wiki pages to export") - format: Literal["markdown", "json"] = Field(..., description="Export format (markdown or json)") - -# --- Model Configuration Models --- -class Model(BaseModel): - """ - Model for LLM model configuration - """ - id: str = Field(..., description="Model identifier") - name: str = Field(..., description="Display name for the model") - -class Provider(BaseModel): - """ - Model for LLM provider configuration - """ - id: str = Field(..., description="Provider identifier") - name: str = Field(..., description="Display name for the provider") - models: List[Model] = Field(..., description="List of available models for this provider") - supportsCustomModel: Optional[bool] = Field(False, description="Whether this provider supports custom models") - -class ModelConfig(BaseModel): - """ - Model for the entire model configuration - """ - providers: List[Provider] = Field(..., description="List of available model providers") - defaultProvider: str = Field(..., description="ID of the default provider") - -class AuthorizationConfig(BaseModel): - code: str = Field(..., description="Authorization code") - -from api.config import configs, WIKI_AUTH_MODE, WIKI_AUTH_CODE @app.get("/lang/config") async def get_lang_config(): @@ -415,9 +319,7 @@ async def read_wiki_cache(owner: str, repo: str, repo_type: str, language: str) cache_path = get_wiki_cache_path(owner, repo, repo_type, language) if os.path.exists(cache_path): try: - with open(cache_path, 'r', encoding='utf-8') as f: - data = json.load(f) - return WikiCacheData(**data) + return await aload(WikiCacheData, cache_path, encoding="utf-8") except Exception as e: logger.error(f"Error reading wiki cache from {cache_path}: {e}") return None @@ -435,18 +337,8 @@ async def save_wiki_cache(data: WikiCacheRequest) -> bool: provider=data.provider, model=data.model ) - # Log size of data to be cached for debugging (avoid logging full content if large) - try: - payload_json = payload.model_dump_json() - payload_size = len(payload_json.encode('utf-8')) - logger.info(f"Payload prepared for caching. Size: {payload_size} bytes.") - except Exception as ser_e: - logger.warning(f"Could not serialize payload for size logging: {ser_e}") - - logger.info(f"Writing cache file to: {cache_path}") - with open(cache_path, 'w', encoding='utf-8') as f: - json.dump(payload.model_dump(), f, indent=2) + await asave(payload, cache_path, encoding="utf-8") logger.info(f"Wiki cache successfully saved to {cache_path}") return True except IOError as e: diff --git a/api/schemas/__init__.py b/api/schemas/__init__.py new file mode 100644 index 000000000..427a9dce3 --- /dev/null +++ b/api/schemas/__init__.py @@ -0,0 +1,36 @@ +from api.schemas.chat import ChatCompletionRequest +from api.schemas.io import aload, asave +from api.schemas.models import ( + AuthorizationConfig, + Model, + ModelConfig, + Provider, +) +from api.schemas.wiki import ( + ProcessedProjectEntry, + RepoInfo, + WikiCacheData, + WikiCacheRequest, + WikiExportRequest, + WikiPage, + WikiSection, + WikiStructureModel, +) + +__all__ = [ + "AuthorizationConfig", + "ChatCompletionRequest", + "Model", + "ModelConfig", + "ProcessedProjectEntry", + "Provider", + "RepoInfo", + "WikiCacheData", + "WikiCacheRequest", + "WikiExportRequest", + "WikiPage", + "WikiSection", + "WikiStructureModel", + "aload", + "asave", +] \ No newline at end of file diff --git a/api/chat_model.py b/api/schemas/chat.py similarity index 60% rename from api/chat_model.py rename to api/schemas/chat.py index 4af8a03d8..d453c6403 100644 --- a/api/chat_model.py +++ b/api/schemas/chat.py @@ -1,32 +1,46 @@ -from typing import Optional, List, Literal +from typing import Literal from urllib.parse import unquote from pydantic import BaseModel, Field, field_validator + # Models for the API class ChatMessage(BaseModel): role: str # 'user' or 'assistant' content: str mode: Literal["normal", "deep_research"] = Field(default="normal") + class ChatCompletionRequest(BaseModel): """ Model for requesting a chat completion. """ + repo_url: str = Field(..., description="URL of the repository to query") - messages: List[ChatMessage] = Field(..., description="List of chat messages") - filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt") - token: Optional[str] = Field(None, description="Personal access token for private repositories") - type: Optional[Literal["local", "github", "gitlab", "bitbucket"]] = Field( + messages: list[ChatMessage] = Field(..., description="List of chat messages") + filePath: str | None = Field( + None, + description="Optional path to a file in the repository to include in the prompt", + ) + token: str | None = Field( + None, description="Personal access token for private repositories" + ) + type: Literal["local", "github", "gitlab", "bitbucket"] | None = Field( "github", - description="Type of repository (e.g., 'local', 'github', 'gitlab', 'bitbucket')", + description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')", ) # model parameters - provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)") - model: Optional[str] = Field(None, description="Model name for the specified provider") + provider: str = Field( + "google", + description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)", + ) + model: str | None = Field(None, description="Model name for the specified provider") - language: Optional[str] = Field("en", description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')") + language: str | None = Field( + "en", + description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')", + ) research_iteration: int = Field( default=1, @@ -34,19 +48,19 @@ class ChatCompletionRequest(BaseModel): description="Current deep research iteration (1-based). Only used when the request is in deep_research mode.", ) - excluded_dirs: List[str] = Field( + excluded_dirs: list[str] = Field( default_factory=list, description="List or newline-separated string of directories to exclude from processing", ) - excluded_files: List[str] = Field( + excluded_files: list[str] = Field( default_factory=list, description="List or newline-separated string of file patterns to exclude from processing", ) - included_dirs: List[str] = Field( + included_dirs: list[str] = Field( default_factory=list, description="List or newline-separated string of directories to include exclusively", ) - included_files: List[str] = Field( + included_files: list[str] = Field( default_factory=list, description="List or newline-separated string of file patterns to include exclusively", ) diff --git a/api/schemas/io.py b/api/schemas/io.py new file mode 100644 index 000000000..c6c64e1e0 --- /dev/null +++ b/api/schemas/io.py @@ -0,0 +1,20 @@ +from typing import TypeVar + +import anyio +from pydantic import BaseModel + +_M = TypeVar("_M", bound=BaseModel) + + +async def asave(model: BaseModel, path: str, *, encoding: str = "utf-8"): + """Asynchronous serialize and save a model""" + + async with await anyio.open_file(path, mode="w", encoding=encoding) as file: + await file.write(model.model_dump_json()) + + +async def aload(model: type[_M], path: str, *, encoding: str = "utf-8") -> _M: + """Asynchronous deserialize and load a model""" + import json + async with await anyio.open_file(path, mode="r", encoding=encoding) as file: + return model(**json.loads(await file.read())) diff --git a/api/schemas/models.py b/api/schemas/models.py new file mode 100644 index 000000000..74a7db912 --- /dev/null +++ b/api/schemas/models.py @@ -0,0 +1,31 @@ +from pydantic import BaseModel, Field + + +class Model(BaseModel): + """ + Model for LLM model configuration + """ + id: str = Field(..., description="Model identifier") + name: str = Field(..., description="Display name for the model") + + +class Provider(BaseModel): + """ + Model for LLM provider configuration + """ + id: str = Field(..., description="Provider identifier") + name: str = Field(..., description="Display name for the provider") + models: list[Model] = Field(..., description="List of available models for this provider") + supportsCustomModel: bool = Field(False, description="Whether this provider supports custom models") + + +class ModelConfig(BaseModel): + """ + Model for the entire model configuration + """ + providers: list[Provider] = Field(..., description="List of available model providers") + defaultProvider: str = Field(..., description="ID of the default provider") + + +class AuthorizationConfig(BaseModel): + code: str = Field(..., description="Authorization code") diff --git a/api/schemas/wiki.py b/api/schemas/wiki.py new file mode 100644 index 000000000..ae6cc0199 --- /dev/null +++ b/api/schemas/wiki.py @@ -0,0 +1,108 @@ +import json +from typing import Literal + +import anyio +from pydantic import BaseModel, Field + + +class RepoInfo(BaseModel): + owner: str + repo: str + type: str + token: str | None = None + localPath: str | None = None + repoUrl: str | None = None + + +class WikiPage(BaseModel): + """ + Model for a wiki page. + """ + + id: str + title: str + content: str + filePaths: list[str] + importance: str # Should ideally be Literal['high', 'medium', 'low'] + relatedPages: list[str] + + +class WikiSection(BaseModel): + """ + Model for the wiki sections. + """ + + id: str + title: str + pages: list[str] + subsections: list[str] | None = None + + +class WikiStructureModel(BaseModel): + """ + Model for the overall wiki structure. + """ + + id: str + title: str + description: str + pages: list[WikiPage] + sections: list[WikiSection] | None = None + rootSections: list[str] | None = None + + +class WikiCacheData(BaseModel): + """ + Model for the data to be stored in the wiki cache. + """ + + wiki_structure: WikiStructureModel + generated_pages: dict[str, WikiPage] + repo_url: str | None = None # compatible for old cache + repo: RepoInfo | None = None + provider: str | None = None + model: str | None = None + + async def save(self, path): + async with await anyio.open_file(path, mode="w", encoding="utf-8") as file: + await file.write(self.model_dump_json()) + + @classmethod + async def load(cls, path): + async with await anyio.open_file(path, mode="r", encoding="utf-8") as file: + return cls(**json.loads(await file.read())) + + +class WikiCacheRequest(BaseModel): + """ + Model for the request body when saving wiki cache. + """ + + repo: RepoInfo + language: str + wiki_structure: WikiStructureModel + generated_pages: dict[str, WikiPage] + provider: str + model: str + + +class WikiExportRequest(BaseModel): + """ + Model for requesting a wiki export. + """ + + repo_url: str = Field(..., description="URL of the repository") + pages: list[WikiPage] = Field(..., description="List of wiki pages to export") + format: Literal["markdown", "json"] = Field( + ..., description="Export format (markdown or json)" + ) + + +class ProcessedProjectEntry(BaseModel): + id: str # Filename + owner: str + repo: str + name: str # owner/repo + repo_type: str # Renamed from type to repo_type for clarity with existing models + submittedAt: int # Timestamp + language: str # Extracted from filename diff --git a/api/simple_chat.py b/api/simple_chat.py index a6e7e8aae..cd661b538 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -17,7 +17,7 @@ DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, SIMPLE_CHAT_SYSTEM_PROMPT ) -from api.chat_model import ChatCompletionRequest +from api.schemas import ChatCompletionRequest # Configure logging from api.logging_config import setup_logging diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index d72efc2b5..6b6949cf9 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -18,7 +18,7 @@ DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, SIMPLE_CHAT_SYSTEM_PROMPT, ) -from api.chat_model import ChatCompletionRequest +from api.schemas import ChatCompletionRequest # Configure logging from api.logging_config import setup_logging From ab41b216ca3e5739a9233542af2246cc0870ee3b Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Tue, 28 Jul 2026 10:08:41 +0800 Subject: [PATCH 2/4] add unittests --- tests/backend/schemas/test_wiki.py | 56 ++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 tests/backend/schemas/test_wiki.py diff --git a/tests/backend/schemas/test_wiki.py b/tests/backend/schemas/test_wiki.py new file mode 100644 index 000000000..93a964387 --- /dev/null +++ b/tests/backend/schemas/test_wiki.py @@ -0,0 +1,56 @@ +import pytest + +from api.schemas import ( + WikiCacheData, + WikiPage, + WikiStructureModel, + aload, + asave, +) + + +@pytest.mark.asyncio +async def test_wiki_cache_async_save(tmp_path) -> None: + cache = WikiCacheData( + wiki_structure=WikiStructureModel( + id="test", title="test", description="test", pages=[] + ), + generated_pages={ + "test": WikiPage( + id="test", + title="test", + content="test", + filePaths=[], + importance="high", + relatedPages=[], + ) + }, + ) + + path = tmp_path / "test.wiki" + await asave(cache, path.as_posix(), encoding="utf-8") + assert path.exists() + +@pytest.mark.asyncio +async def test_wiki_cache_async_load(tmp_path) -> None: + cache = WikiCacheData( + wiki_structure=WikiStructureModel( + id="test", title="test", description="test", pages=[] + ), + generated_pages={ + "test": WikiPage( + id="test", + title="test", + content="test", + filePaths=[], + importance="high", + relatedPages=[], + ) + }, + ) + + path = tmp_path / "test.wiki" + await asave(cache, path.as_posix(), encoding="utf-8") + ret = await aload(WikiCacheData, path.as_posix(), encoding="utf-8") + + assert ret == cache From cc5efec8fa984043a1f89d7a7da13998fe61b6ed Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Tue, 28 Jul 2026 18:48:31 +0800 Subject: [PATCH 3/4] ruff formating --- api/schemas/__init__.py | 2 +- api/schemas/io.py | 1 + api/schemas/models.py | 15 ++++++++++++--- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/api/schemas/__init__.py b/api/schemas/__init__.py index 427a9dce3..a18a9e876 100644 --- a/api/schemas/__init__.py +++ b/api/schemas/__init__.py @@ -33,4 +33,4 @@ "WikiStructureModel", "aload", "asave", -] \ No newline at end of file +] diff --git a/api/schemas/io.py b/api/schemas/io.py index c6c64e1e0..4e215a48b 100644 --- a/api/schemas/io.py +++ b/api/schemas/io.py @@ -16,5 +16,6 @@ async def asave(model: BaseModel, path: str, *, encoding: str = "utf-8"): async def aload(model: type[_M], path: str, *, encoding: str = "utf-8") -> _M: """Asynchronous deserialize and load a model""" import json + async with await anyio.open_file(path, mode="r", encoding=encoding) as file: return model(**json.loads(await file.read())) diff --git a/api/schemas/models.py b/api/schemas/models.py index 74a7db912..7cf5fe6ee 100644 --- a/api/schemas/models.py +++ b/api/schemas/models.py @@ -5,6 +5,7 @@ class Model(BaseModel): """ Model for LLM model configuration """ + id: str = Field(..., description="Model identifier") name: str = Field(..., description="Display name for the model") @@ -13,17 +14,25 @@ class Provider(BaseModel): """ Model for LLM provider configuration """ + id: str = Field(..., description="Provider identifier") name: str = Field(..., description="Display name for the provider") - models: list[Model] = Field(..., description="List of available models for this provider") - supportsCustomModel: bool = Field(False, description="Whether this provider supports custom models") + models: list[Model] = Field( + ..., description="List of available models for this provider" + ) + supportsCustomModel: bool = Field( + False, description="Whether this provider supports custom models" + ) class ModelConfig(BaseModel): """ Model for the entire model configuration """ - providers: list[Provider] = Field(..., description="List of available model providers") + + providers: list[Provider] = Field( + ..., description="List of available model providers" + ) defaultProvider: str = Field(..., description="ID of the default provider") From 98418a3c955414d53d83ed26c0afed736f2c3115 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 29 Jul 2026 18:11:44 +0800 Subject: [PATCH 4/4] Apply suggestion from @GdoongMathew --- api/schemas/wiki.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/api/schemas/wiki.py b/api/schemas/wiki.py index ae6cc0199..12c1da2c0 100644 --- a/api/schemas/wiki.py +++ b/api/schemas/wiki.py @@ -63,15 +63,6 @@ class WikiCacheData(BaseModel): provider: str | None = None model: str | None = None - async def save(self, path): - async with await anyio.open_file(path, mode="w", encoding="utf-8") as file: - await file.write(self.model_dump_json()) - - @classmethod - async def load(cls, path): - async with await anyio.open_file(path, mode="r", encoding="utf-8") as file: - return cls(**json.loads(await file.read())) - class WikiCacheRequest(BaseModel): """