modified: app.py

modified: requirements.txt
modified: src/ai_infra.py
modified: src/ai_transform.py

app.py

@@ -4,6 +4,7 @@ from src.ai_infra import init_ai_config, get_ai_models
 from src.utils import *
 import requests, os, datetime
 import json
+from typing import Any
 
 # for offline debugging
 if os.path.exists("secret.env"):
@@ -19,7 +20,7 @@ NOTION_TOKEN = os.environ.get("NOTION_TOKEN")
 DB_ID = os.environ.get("DB_ID")
 
 def get_property(it, task_name, url, api_key, model):
-    properties = {
+    properties: dict[str, Any] = {
         "Task name": {
             "title": [{"text": {"content": task_name}}]
         }
@@ -72,6 +73,8 @@ def init_items(items):
 
 @app.route("/", methods=["GET"])
 def index():
+    if app.static_folder is None:
+        raise RuntimeError("Static folder is not configured")
     return send_from_directory(app.static_folder, "index.html")
 
 # @app.route('/', methods=['GET'])

requirements.txt

@@ -1,4 +1,5 @@
 flask==2.3.3
 requests==2.31.0
 python-dotenv
-openai>=0.27.8
+openai>=0.27.8
+google-genai==1.50.0

src/ai_infra.py

@@ -1,8 +1,15 @@
 
 import os
 from openai import OpenAI
+from openai.types.chat import (
+    ChatCompletionUserMessageParam,
+    ChatCompletionSystemMessageParam,
+)
+from google import genai
+from typing import Union
+import io
 
-def chat_completion(question: str, model: str, base_url: str, api_key: str, system_instr: str = None) -> str:
+def chat_completion(question: str, model: str, base_url: str, api_key: str, system_instr: str | None = None) -> str:
     '''
     ai interaction function using OpenAI SDK.
     Parameters
@@ -20,21 +27,59 @@ def chat_completion(question: str, model: str, base_url: str, api_key: str, syst
     '''
     client = OpenAI(api_key=api_key, 
                     base_url=base_url)
-    messages = [{"role": "user", "content": question}]
+    messages: list[ChatCompletionUserMessageParam |
+                ChatCompletionSystemMessageParam] = [
+        ChatCompletionUserMessageParam(role="user", content=question)
+    ]
     if system_instr is not None:
-        messages.append({
-            "role": "system",
-            "content": system_instr
-        })
+        messages.append(
+            ChatCompletionSystemMessageParam(role="system", content=system_instr)
+        )
     response = client.chat.completions.create(
         # reasoning_effort="high",
         model=model,
         messages=messages
     )
-    content = response.choices[0].message.content
+    content = response.choices[0].message.content or ""
     print(f"Model: {model}, Base URL: {base_url}\nResponse: {content}\n")
     return content
 
+def init_genai_client(api_key: str | None = None) -> genai.Client:
+    if not api_key:
+        api_key = os.environ.get(f"GEMINI_KEY") or ""
+    if not api_key:
+        raise ValueError("API key for Gemini is not provided.")
+    return genai.Client(api_key=api_key)
+
+def interact_with_pdf(client: genai.Client, file: Union[str, os.PathLike[str], io.IOBase], question: str = "") -> str:
+    '''
+    ai interaction function using Generative SDK to interact with PDF files.
+    Parameters
+    ----------
+    file : 
+        A path to the file or an `IOBase` object to be uploaded. If it's an
+        IOBase object, it must be opened in blocking (the default) mode and
+        binary mode. In other words, do not use non-blocking mode or text mode.
+        The given stream must be seekable, that is, it must be able to call
+        `seek()` on 'path'.
+    question : str
+        The input question or prompt to send to the AI model.
+    '''
+
+    # 上传文件
+    pdf_file = client.files.upload(file=file)
+    # 生成内容
+    response = client.models.generate_content(
+        model="gemini-2.5-flash",
+        contents=[
+            question,
+            pdf_file
+        ]
+    )
+    content = response.text or ""
+    print(content)
+    return content
+
 def init_ai_config(model: str = "default")-> dict[str, str]:
     '''
     Initialize configuration for different AI models based on the model name.
@@ -110,7 +155,7 @@ def init_ai_config(model: str = "default")-> dict[str, str]:
         }
         API_MODEL = "GPT"
     # config["api_key"] = "Bearer " + os.environ.get(f"{API_MODEL}_KEY")
-    config["api_key"] = os.environ.get(f"{API_MODEL}_KEY")
+    config["api_key"] = os.environ.get(f"{API_MODEL}_KEY") or ""
     return config
 
 def get_ai_models() -> list:

src/ai_transform.py

@@ -9,7 +9,7 @@ from typing import Union, List
 with open("prompt.txt", encoding="utf-8") as f:
     system_content = f.read()
 
-def modified_with_ai(items, url, api_key, model=None):
+def modified_with_ai(items, url, api_key, model):
     # 在这里调用AI模型对item进行修改
     # today_date = time.strftime("%Y-%m-%d", time.localtime())
     today_date = time_cali()
@@ -49,10 +49,13 @@ def classify_task_with_ai(
         返回任务所属的类别名称，或在不匹配任何类别时返回 ''。
     """
     # 1. 读取并格式化任务类别 (RAG中的 "R" - Retrieval)
+    if config is None and model is None:
+        raise ValueError("Either config or model must be provided.")
     if config is not None:
         url = config["url"]
         api_key = config["api_key"]
         model = config["model"]
+    assert isinstance(model, str)
     categories_mapping = load_task_mapping_from_txt(categories_filepath)
     task_list = categories_mapping.keys()
     # 2. 定义分类 prompt 模板