mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-05-09 05:42:39 +08:00
Compare commits
434 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9620a06552 | ||
|
|
9b00a5f3f1 | ||
|
|
faa77be843 | ||
|
|
28f158c479 | ||
|
|
90c3afcfa4 | ||
|
|
565e10b6a5 | ||
|
|
773ed5e6f7 | ||
|
|
8351312b2b | ||
|
|
41f53d39a0 | ||
|
|
4873ffda84 | ||
|
|
b79609bb8b | ||
|
|
bdcbb5cce6 | ||
|
|
d1503f9df3 | ||
|
|
210c3234d2 | ||
|
|
c13abfdd0d | ||
|
|
30b332ac7e | ||
|
|
7e9c489aeb | ||
|
|
5739ca7f97 | ||
|
|
e4451c7e6a | ||
|
|
5cded77387 | ||
|
|
ea4e0dd764 | ||
|
|
f105357f96 | ||
|
|
bc2302baeb | ||
|
|
afcdefbbf3 | ||
|
|
3ad8557065 | ||
|
|
e68d607c9b | ||
|
|
8e9cf67190 | ||
|
|
0cb6cd8761 | ||
|
|
17aa795b3e | ||
|
|
7d47096e6e | ||
|
|
48b59df11b | ||
|
|
a90a3b2445 | ||
|
|
d18b68d24a | ||
|
|
78c4ec8bfe | ||
|
|
b50a3b9aae | ||
|
|
4f3eaa12d5 | ||
|
|
cedb0f565c | ||
|
|
226432ec7f | ||
|
|
d93ab0143c | ||
|
|
3d32d66ab1 | ||
|
|
e814eed047 | ||
|
|
96395c1469 | ||
|
|
6065c29891 | ||
|
|
f38cb274e4 | ||
|
|
7bfee87cbf | ||
|
|
2ce2a3754c | ||
|
|
510476c214 | ||
|
|
6cd071c84b | ||
|
|
406e17b3fa | ||
|
|
dd184255ad | ||
|
|
77a0b38081 | ||
|
|
14c3d66ce6 | ||
|
|
858da38680 | ||
|
|
9f381b3c73 | ||
|
|
b8fc20b981 | ||
|
|
b89825525a | ||
|
|
e09cfc6704 | ||
|
|
0c9c303c60 | ||
|
|
3156b43739 | ||
|
|
591aa990a6 | ||
|
|
3be29f36a7 | ||
|
|
7638db4c3b | ||
|
|
0312a500a6 | ||
|
|
1a88b5355a | ||
|
|
3374773de5 | ||
|
|
872b5fe3da | ||
|
|
be15e9871c | ||
|
|
024a6a253b | ||
|
|
1af662df7b | ||
|
|
b4f64eb593 | ||
|
|
86aa86208c | ||
|
|
018e814615 | ||
|
|
e4d6e5cfc7 | ||
|
|
770cd77632 | ||
|
|
9f1692b33d | ||
|
|
6f63e0a5d7 | ||
|
|
6a90e2c796 | ||
|
|
23b90ff0f9 | ||
|
|
dc86af2fa4 | ||
|
|
425b822046 | ||
|
|
65c18b1d52 | ||
|
|
1bddf3daa7 | ||
|
|
600b6af876 | ||
|
|
4bdf16331d | ||
|
|
87cbda0528 | ||
|
|
9897941bf9 | ||
|
|
31938812d0 | ||
|
|
19d879d3f6 | ||
|
|
cc41036c63 | ||
|
|
a9f2b40529 | ||
|
|
86000ea19a | ||
|
|
0422c3b9e7 | ||
|
|
64c8bd5b5a | ||
|
|
a7eba2c5fc | ||
|
|
2b7753e43e | ||
|
|
47c1e5b5b8 | ||
|
|
14ee97def0 | ||
|
|
92e262f732 | ||
|
|
c46880b701 | ||
|
|
473e9b9300 | ||
|
|
28945ef153 | ||
|
|
b6b5d9f9c4 | ||
|
|
ba5de1ab31 | ||
|
|
002ebeaade | ||
|
|
894756000c | ||
|
|
cdb178c503 | ||
|
|
7c48cafc71 | ||
|
|
74d4592238 | ||
|
|
0044dd104e | ||
|
|
05041e2eae | ||
|
|
78908f216d | ||
|
|
efc68ae701 | ||
|
|
e9340a8b4b | ||
|
|
66e199d516 | ||
|
|
6151d8a787 | ||
|
|
296261da8a | ||
|
|
383371dd6f | ||
|
|
bb8c026bda | ||
|
|
344993dd6f | ||
|
|
ffb048c314 | ||
|
|
3eef9b8faa | ||
|
|
5704bb646b | ||
|
|
fbc684b3a7 | ||
|
|
6529b2a9c3 | ||
|
|
a1701e2edf | ||
|
|
eba6391de7 | ||
|
|
9f2c3c9688 | ||
|
|
57f5a19d0c | ||
|
|
c8d53c6964 | ||
|
|
643cda1abe | ||
|
|
03d118a73a | ||
|
|
51dd7f5c17 | ||
|
|
af7e1e7a3c | ||
|
|
ea5d855bc3 | ||
|
|
5f74367cd6 | ||
|
|
26e41e1c14 | ||
|
|
1bb2b50043 | ||
|
|
7bdb629f03 | ||
|
|
fd92f986da | ||
|
|
69a1207102 | ||
|
|
def652c768 | ||
|
|
c35faf5356 | ||
|
|
0615a33206 | ||
|
|
e77530bdc5 | ||
|
|
8c62df63cc | ||
|
|
bd36eade77 | ||
|
|
d2c023081a | ||
|
|
63d0850b38 | ||
|
|
c86659428f | ||
|
|
bf7cc6caf0 | ||
|
|
26b8be6041 | ||
|
|
f978f9196f | ||
|
|
75cb8d2a3c | ||
|
|
17a21ed707 | ||
|
|
f390647139 | ||
|
|
aacd91e196 | ||
|
|
258171c9c4 | ||
|
|
812c5873aa | ||
|
|
4c3d47f1f0 | ||
|
|
ba7b6ba869 | ||
|
|
d0471ae512 | ||
|
|
636c4be9fb | ||
|
|
6bec765a9d | ||
|
|
d61d16ccc4 | ||
|
|
f2a5715b24 | ||
|
|
c064c3781f | ||
|
|
bb4dffe2a4 | ||
|
|
37cf3eeef3 | ||
|
|
40395b2999 | ||
|
|
32afe6445f | ||
|
|
793a991913 | ||
|
|
d278224ff1 | ||
|
|
9b4d0ce6a8 | ||
|
|
a1829fe590 | ||
|
|
2b2b39365c | ||
|
|
1147930f3f | ||
|
|
636f338ed7 | ||
|
|
72365d00b4 | ||
|
|
19d8086732 | ||
|
|
30488418e5 | ||
|
|
2f0badd74a | ||
|
|
6045b0579b | ||
|
|
498f1fec74 | ||
|
|
f6a541f2b9 | ||
|
|
8ce78eabca | ||
|
|
2c34c5309f | ||
|
|
77e680168a | ||
|
|
8a7e59742f | ||
|
|
42bac14770 | ||
|
|
8323834483 | ||
|
|
1751caef62 | ||
|
|
d622d1474d | ||
|
|
f28be2e7de | ||
|
|
17773913ae | ||
|
|
d469c2d3f9 | ||
|
|
4e74d32882 | ||
|
|
7b8cd37a9b | ||
|
|
eda306d726 | ||
|
|
94f3b1fe84 | ||
|
|
c50e3ba293 | ||
|
|
eff7818912 | ||
|
|
270bcff8f3 | ||
|
|
e04963c2dc | ||
|
|
f369967c91 | ||
|
|
cd982c5526 | ||
|
|
16e03c9d37 | ||
|
|
d38b1f5364 | ||
|
|
f57ba4d05e | ||
|
|
172eeaafcf | ||
|
|
3115ed28b2 | ||
|
|
d8dc53805c | ||
|
|
7218d10e1b | ||
|
|
89bf85f501 | ||
|
|
8334a468d0 | ||
|
|
3da80ed077 | ||
|
|
2883ccbe87 | ||
|
|
5d3443fee4 | ||
|
|
27756a53db | ||
|
|
71cde6661d | ||
|
|
a857337b31 | ||
|
|
4ee21ffae4 | ||
|
|
d8399f7e85 | ||
|
|
574ac8d32f | ||
|
|
a2611bfa7d | ||
|
|
853badb76f | ||
|
|
5d69e1d2a5 | ||
|
|
6494f28bdb | ||
|
|
f55916bda2 | ||
|
|
04691ee197 | ||
|
|
2ac0e564e1 | ||
|
|
6072a29a20 | ||
|
|
8658942385 | ||
|
|
cc4859950c | ||
|
|
23b81ad6f1 | ||
|
|
e3b9dca5c0 | ||
|
|
a2359a1ad2 | ||
|
|
cb875b1b34 | ||
|
|
b92a85b4bc | ||
|
|
8c7dd6bab2 | ||
|
|
aad7df64d7 | ||
|
|
8474342007 | ||
|
|
61ccb4be65 | ||
|
|
1c6f69707c | ||
|
|
e08e8c482a | ||
|
|
548c1d2cab | ||
|
|
5a071bf3d1 | ||
|
|
1bffcbd947 | ||
|
|
274a36a83a | ||
|
|
ec40f36114 | ||
|
|
af19f274a7 | ||
|
|
2316004194 | ||
|
|
98762198ef | ||
|
|
1469de22a4 | ||
|
|
1e687f960a | ||
|
|
7f01b835fd | ||
|
|
e46b6c5c01 | ||
|
|
74226ad8df | ||
|
|
f8ae7be539 | ||
|
|
37b16e380d | ||
|
|
9ea3e9f652 | ||
|
|
54422b5181 | ||
|
|
712995dcf3 | ||
|
|
c2767b0fd6 | ||
|
|
179cc61f65 | ||
|
|
f3b910d55a | ||
|
|
f4157b52ea | ||
|
|
79710310ce | ||
|
|
3412498438 | ||
|
|
b896b07a08 | ||
|
|
379bff0622 | ||
|
|
474f47aa9f | ||
|
|
f1e26a4133 | ||
|
|
e37f881207 | ||
|
|
306c0b707b | ||
|
|
08c448ee30 | ||
|
|
1532014067 | ||
|
|
fa9f604af9 | ||
|
|
3b3d0d6539 | ||
|
|
9641d33040 | ||
|
|
eca339d107 | ||
|
|
ca18705d88 | ||
|
|
8f17b52466 | ||
|
|
8cf84e722b | ||
|
|
7c4d736b54 | ||
|
|
1b3ae6ab25 | ||
|
|
a4ad08136e | ||
|
|
df5e7997c5 | ||
|
|
b2cb3768c1 | ||
|
|
fa169c5cd3 | ||
|
|
bbb3975b67 | ||
|
|
4502a9c4fa | ||
|
|
86905a2670 | ||
|
|
b1e60a4867 | ||
|
|
1efe3324fb | ||
|
|
55c1e37d39 | ||
|
|
7fa700317c | ||
|
|
bbe831a57c | ||
|
|
90c86c056c | ||
|
|
36f22a28df | ||
|
|
ac03c51e2c | ||
|
|
bd9e92f705 | ||
|
|
281eff5eb2 | ||
|
|
abbd2253ad | ||
|
|
46466624ae | ||
|
|
0ba8d51b2a | ||
|
|
a1408ee18f | ||
|
|
58030bbcff | ||
|
|
e1b3e6ef01 | ||
|
|
298a6ba8ab | ||
|
|
e5bf47629f | ||
|
|
ea29ee9f66 | ||
|
|
868c2254de | ||
|
|
567522c87a | ||
|
|
25fd47f57b | ||
|
|
f89d6342d1 | ||
|
|
b02affdea3 | ||
|
|
6e5ade943b | ||
|
|
a6ed0c0d00 | ||
|
|
68402aadd7 | ||
|
|
85cacd447b | ||
|
|
11262b321a | ||
|
|
bf290f063d | ||
|
|
7ac0fbaf76 | ||
|
|
7489c76722 | ||
|
|
bcdf1b6efe | ||
|
|
8a9dbe212c | ||
|
|
16bd71a6cb | ||
|
|
71caad0655 | ||
|
|
2c62ffe34a | ||
|
|
3450a89880 | ||
|
|
a081a69bbe | ||
|
|
271d1d23d5 | ||
|
|
605aba1a3c | ||
|
|
be3c2b4c7c | ||
|
|
08eb32d7bd | ||
|
|
2b9cda15e4 | ||
|
|
f6055b290a | ||
|
|
ec665e05e4 | ||
|
|
2b6d7205ec | ||
|
|
41381a920c | ||
|
|
f1b3fc2254 | ||
|
|
a677ed307d | ||
|
|
0ab23ee972 | ||
|
|
43f56d39be | ||
|
|
a39caee5f5 | ||
|
|
2edfdf47c8 | ||
|
|
3819461db5 | ||
|
|
85654dd7dd | ||
|
|
619a70416b | ||
|
|
16d996fe70 | ||
|
|
1baeb6da19 | ||
|
|
1641d432dd | ||
|
|
1bf9862e47 | ||
|
|
602a394043 | ||
|
|
22a2415ca5 | ||
|
|
feb034352d | ||
|
|
a7c8942c78 | ||
|
|
95f2ac3811 | ||
|
|
91354295f2 | ||
|
|
c9c4ab5911 | ||
|
|
a26c5e40dd | ||
|
|
80f5c7bc44 | ||
|
|
4833b39c52 | ||
|
|
f478958943 | ||
|
|
0469ad46d6 | ||
|
|
5fe5deb9df | ||
|
|
ce83bc24bd | ||
|
|
dce729c8cb | ||
|
|
a9d17cd96f | ||
|
|
294bb3d4a1 | ||
|
|
b31b9261f2 | ||
|
|
2211f8d9e4 | ||
|
|
b9b7b00a7f | ||
|
|
843faf6103 | ||
|
|
4af5dad9a8 | ||
|
|
52437c9d18 | ||
|
|
c6cb4c8479 | ||
|
|
c3714ec251 | ||
|
|
dbe2f94af1 | ||
|
|
07fd5f8a9e | ||
|
|
9e64b4cd7f | ||
|
|
f08a7b9eb3 | ||
|
|
a6fa764e2a | ||
|
|
01676668f1 | ||
|
|
8e5e4f460d | ||
|
|
f907b8a84d | ||
|
|
a3a4285f90 | ||
|
|
0979163b79 | ||
|
|
248a25eaee | ||
|
|
f95b1fa68a | ||
|
|
d2b5d69051 | ||
|
|
3ca419b735 | ||
|
|
50e275a2f9 | ||
|
|
aeccf78957 | ||
|
|
cb3cef70e5 | ||
|
|
b9bd303bf8 | ||
|
|
57d4786a7f | ||
|
|
df031455b2 | ||
|
|
30059eff4f | ||
|
|
bc289b48c8 | ||
|
|
067d8b99b8 | ||
|
|
00a6a9c42d | ||
|
|
070425d446 | ||
|
|
7405883444 | ||
|
|
66959937ed | ||
|
|
e431efbcba | ||
|
|
ba00baa5a0 | ||
|
|
0fb5d4a164 | ||
|
|
1ac717b67f | ||
|
|
273cbd447e | ||
|
|
cee41567a2 | ||
|
|
1aae5eb1a6 | ||
|
|
28a4c81aff | ||
|
|
5e077cd64d | ||
|
|
e3f957a59b | ||
|
|
55c62a3ab5 | ||
|
|
22e7eef1bd | ||
|
|
d6524907f3 | ||
|
|
357db334cd | ||
|
|
f8bed3909b | ||
|
|
182bbdde91 | ||
|
|
2c70f990c2 | ||
|
|
0b01a6aa91 | ||
|
|
e557dffbc6 | ||
|
|
7f33b0b1b8 | ||
|
|
41ddf77a5b | ||
|
|
8c657ce41d | ||
|
|
3ff3b9ed4a | ||
|
|
ef43419ecd | ||
|
|
2ca375c214 | ||
|
|
cbd45c1d0f | ||
|
|
2592ea3464 | ||
|
|
73ac97cd96 | ||
|
|
e014663e97 |
1
.github/workflows/issues.yml
vendored
1
.github/workflows/issues.yml
vendored
@@ -29,4 +29,5 @@ jobs:
|
||||
days-before-pr-close: -1
|
||||
# 排除带有RFC标签的issue
|
||||
exempt-issue-labels: "RFC"
|
||||
operations-per-run: 500
|
||||
repo-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
5
.gitignore
vendored
5
.gitignore
vendored
@@ -27,4 +27,7 @@ venv
|
||||
|
||||
# Pylint
|
||||
pylint-report.json
|
||||
.pylint.d/
|
||||
.pylint.d/
|
||||
|
||||
# AI
|
||||
.claude/
|
||||
|
||||
@@ -26,6 +26,11 @@
|
||||
|
||||
官方Wiki:https://wiki.movie-pilot.org
|
||||
|
||||
### 为 AI Agent 添加 Skills
|
||||
```shell
|
||||
npx skills add https://github.com/jxxghp/MoviePilot
|
||||
```
|
||||
|
||||
## 参与开发
|
||||
|
||||
API文档:https://api.movie-pilot.org
|
||||
|
||||
@@ -1,22 +1,27 @@
|
||||
"""MoviePilot AI智能体实现"""
|
||||
|
||||
import asyncio
|
||||
from typing import Dict, List, Any
|
||||
import traceback
|
||||
from time import strftime
|
||||
from typing import Dict, List
|
||||
|
||||
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
||||
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
|
||||
from langchain_community.callbacks import get_openai_callback
|
||||
from langchain_core.chat_history import InMemoryChatMessageHistory
|
||||
from langchain_core.messages import HumanMessage, AIMessage, ToolCall, ToolMessage, SystemMessage
|
||||
from langchain_core.runnables.history import RunnableWithMessageHistory
|
||||
from langchain.agents import create_agent
|
||||
from langchain.agents.middleware import (
|
||||
SummarizationMiddleware, LLMToolSelectorMiddleware,
|
||||
)
|
||||
from langchain_core.messages import (
|
||||
HumanMessage,
|
||||
BaseMessage,
|
||||
)
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from app.agent.callback import StreamingCallbackHandler
|
||||
from app.agent.memory import ConversationMemoryManager
|
||||
from app.agent.prompt import PromptManager
|
||||
from app.agent.callback import StreamingHandler
|
||||
from app.agent.memory import memory_manager
|
||||
from app.agent.middleware.memory import MemoryMiddleware
|
||||
from app.agent.middleware.patch_tool_calls import PatchToolCallsMiddleware
|
||||
from app.agent.prompt import prompt_manager
|
||||
from app.agent.tools.factory import MoviePilotToolFactory
|
||||
from app.chain import ChainBase
|
||||
from app.core.config import settings
|
||||
from app.helper.message import MessageHelper
|
||||
from app.helper.llm import LLMHelper
|
||||
from app.log import logger
|
||||
from app.schemas import Notification
|
||||
|
||||
@@ -26,273 +31,175 @@ class AgentChain(ChainBase):
|
||||
|
||||
|
||||
class MoviePilotAgent:
|
||||
"""MoviePilot AI智能体"""
|
||||
"""
|
||||
MoviePilot AI智能体(基于 LangChain v1 + LangGraph)
|
||||
"""
|
||||
|
||||
def __init__(self, session_id: str, user_id: str = None,
|
||||
channel: str = None, source: str = None, username: str = None):
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
user_id: str = None,
|
||||
channel: str = None,
|
||||
source: str = None,
|
||||
username: str = None,
|
||||
):
|
||||
self.session_id = session_id
|
||||
self.user_id = user_id
|
||||
self.channel = channel # 消息渠道
|
||||
self.source = source # 消息来源
|
||||
self.username = username # 用户名
|
||||
self.channel = channel
|
||||
self.source = source
|
||||
self.username = username
|
||||
|
||||
# 消息助手
|
||||
self.message_helper = MessageHelper()
|
||||
# 流式token管理
|
||||
self.stream_handler = StreamingHandler()
|
||||
|
||||
# 记忆管理器
|
||||
self.memory_manager = ConversationMemoryManager()
|
||||
|
||||
# 提示词管理器
|
||||
self.prompt_manager = PromptManager()
|
||||
|
||||
# 回调处理器
|
||||
self.callback_handler = StreamingCallbackHandler(
|
||||
session_id=session_id
|
||||
)
|
||||
|
||||
# LLM模型
|
||||
self.llm = self._initialize_llm()
|
||||
|
||||
# 工具
|
||||
self.tools = self._initialize_tools()
|
||||
|
||||
# 提示词模板
|
||||
self.prompt = self._initialize_prompt()
|
||||
|
||||
# Agent执行器
|
||||
self.agent_executor = self._create_agent_executor()
|
||||
|
||||
def _initialize_llm(self):
|
||||
"""初始化LLM模型"""
|
||||
provider = settings.LLM_PROVIDER.lower()
|
||||
api_key = settings.LLM_API_KEY
|
||||
|
||||
if provider == "google":
|
||||
if settings.PROXY_HOST:
|
||||
from langchain_openai import ChatOpenAI
|
||||
return ChatOpenAI(
|
||||
model=settings.LLM_MODEL,
|
||||
api_key=api_key,
|
||||
max_retries=3,
|
||||
base_url="https://generativelanguage.googleapis.com/v1beta/openai",
|
||||
temperature=settings.LLM_TEMPERATURE,
|
||||
streaming=True,
|
||||
callbacks=[self.callback_handler],
|
||||
stream_usage=True,
|
||||
openai_proxy=settings.PROXY_HOST
|
||||
)
|
||||
else:
|
||||
from langchain_google_genai import ChatGoogleGenerativeAI
|
||||
return ChatGoogleGenerativeAI(
|
||||
model=settings.LLM_MODEL,
|
||||
google_api_key=api_key,
|
||||
max_retries=3,
|
||||
temperature=settings.LLM_TEMPERATURE,
|
||||
streaming=True,
|
||||
callbacks=[self.callback_handler]
|
||||
)
|
||||
elif provider == "deepseek":
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
return ChatDeepSeek(
|
||||
model=settings.LLM_MODEL,
|
||||
api_key=api_key,
|
||||
max_retries=3,
|
||||
temperature=settings.LLM_TEMPERATURE,
|
||||
streaming=True,
|
||||
callbacks=[self.callback_handler],
|
||||
stream_usage=True
|
||||
)
|
||||
else:
|
||||
from langchain_openai import ChatOpenAI
|
||||
return ChatOpenAI(
|
||||
model=settings.LLM_MODEL,
|
||||
api_key=api_key,
|
||||
max_retries=3,
|
||||
base_url=settings.LLM_BASE_URL,
|
||||
temperature=settings.LLM_TEMPERATURE,
|
||||
streaming=True,
|
||||
callbacks=[self.callback_handler],
|
||||
stream_usage=True,
|
||||
openai_proxy=settings.PROXY_HOST
|
||||
)
|
||||
@staticmethod
|
||||
def _initialize_llm():
|
||||
"""
|
||||
初始化 LLM(带流式回调)
|
||||
"""
|
||||
return LLMHelper.get_llm(streaming=True)
|
||||
|
||||
def _initialize_tools(self) -> List:
|
||||
"""初始化工具列表"""
|
||||
"""
|
||||
初始化工具列表
|
||||
"""
|
||||
return MoviePilotToolFactory.create_tools(
|
||||
session_id=self.session_id,
|
||||
user_id=self.user_id,
|
||||
channel=self.channel,
|
||||
source=self.source,
|
||||
username=self.username,
|
||||
callback_handler=self.callback_handler,
|
||||
memory_mananger=self.memory_manager
|
||||
stream_handler=self.stream_handler,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _initialize_session_store() -> Dict[str, InMemoryChatMessageHistory]:
|
||||
"""初始化内存存储"""
|
||||
return {}
|
||||
|
||||
def get_session_history(self, session_id: str) -> InMemoryChatMessageHistory:
|
||||
"""获取会话历史"""
|
||||
chat_history = InMemoryChatMessageHistory()
|
||||
messages: List[dict] = self.memory_manager.get_recent_messages_for_agent(
|
||||
session_id=session_id,
|
||||
user_id=self.user_id
|
||||
)
|
||||
if messages:
|
||||
for msg in messages:
|
||||
if msg.get("role") == "user":
|
||||
chat_history.add_message(HumanMessage(content=msg.get("content", "")))
|
||||
elif msg.get("role") == "agent":
|
||||
chat_history.add_message(AIMessage(content=msg.get("content", "")))
|
||||
elif msg.get("role") == "tool_call":
|
||||
metadata = msg.get("metadata", {})
|
||||
chat_history.add_message(
|
||||
AIMessage(
|
||||
content=msg.get("content", ""),
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
id=metadata.get("call_id"),
|
||||
name=metadata.get("tool_name"),
|
||||
args=metadata.get("parameters"),
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
elif msg.get("role") == "tool_result":
|
||||
chat_history.add_message(ToolMessage(content=msg.get("content", "")))
|
||||
elif msg.get("role") == "system":
|
||||
chat_history.add_message(SystemMessage(content=msg.get("content", "")))
|
||||
return chat_history
|
||||
|
||||
@staticmethod
|
||||
def _initialize_prompt() -> ChatPromptTemplate:
|
||||
"""初始化提示词模板"""
|
||||
def _create_agent(self):
|
||||
"""
|
||||
创建 LangGraph Agent(使用 create_agent + SummarizationMiddleware)
|
||||
"""
|
||||
try:
|
||||
prompt_template = ChatPromptTemplate.from_messages([
|
||||
("system", "{system_prompt}"),
|
||||
MessagesPlaceholder(variable_name="chat_history"),
|
||||
("user", "{input}"),
|
||||
MessagesPlaceholder(variable_name="agent_scratchpad"),
|
||||
])
|
||||
logger.info("LangChain提示词模板初始化成功")
|
||||
return prompt_template
|
||||
except Exception as e:
|
||||
logger.error(f"初始化提示词失败: {e}")
|
||||
raise e
|
||||
# 系统提示词
|
||||
system_prompt = prompt_manager.get_agent_prompt(
|
||||
channel=self.channel
|
||||
).format(
|
||||
current_date=strftime('%Y-%m-%d')
|
||||
)
|
||||
|
||||
def _create_agent_executor(self) -> RunnableWithMessageHistory:
|
||||
"""创建Agent执行器"""
|
||||
try:
|
||||
agent = create_openai_tools_agent(
|
||||
llm=self.llm,
|
||||
tools=self.tools,
|
||||
prompt=self.prompt
|
||||
)
|
||||
executor = AgentExecutor(
|
||||
agent=agent,
|
||||
tools=self.tools,
|
||||
verbose=settings.LLM_VERBOSE,
|
||||
max_iterations=settings.LLM_MAX_ITERATIONS,
|
||||
return_intermediate_steps=True,
|
||||
handle_parsing_errors=True,
|
||||
early_stopping_method="force"
|
||||
)
|
||||
return RunnableWithMessageHistory(
|
||||
executor,
|
||||
self.get_session_history,
|
||||
input_messages_key="input",
|
||||
history_messages_key="chat_history"
|
||||
# LLM 模型(用于 agent 执行)
|
||||
llm = self._initialize_llm()
|
||||
|
||||
# 工具列表
|
||||
tools = self._initialize_tools()
|
||||
|
||||
# 中间件
|
||||
middlewares = [
|
||||
# 工具选择
|
||||
LLMToolSelectorMiddleware(
|
||||
model=llm,
|
||||
max_tools=20
|
||||
),
|
||||
# 记忆管理
|
||||
MemoryMiddleware(
|
||||
sources=[str(settings.CONFIG_PATH / "agent" / "MEMORY.md")]
|
||||
),
|
||||
# 上下文压缩
|
||||
SummarizationMiddleware(
|
||||
model=llm,
|
||||
trigger=("fraction", 0.85)
|
||||
),
|
||||
# 错误工具调用修复
|
||||
PatchToolCallsMiddleware()
|
||||
]
|
||||
|
||||
return create_agent(
|
||||
model=llm,
|
||||
tools=tools,
|
||||
system_prompt=system_prompt,
|
||||
middleware=middlewares,
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"创建Agent执行器失败: {e}")
|
||||
logger.error(f"创建 Agent 失败: {e}")
|
||||
raise e
|
||||
|
||||
async def process_message(self, message: str) -> str:
|
||||
"""处理用户消息"""
|
||||
async def process(self, message: str) -> str:
|
||||
"""
|
||||
处理用户消息,流式推理并返回 Agent 回复
|
||||
"""
|
||||
try:
|
||||
# 添加用户消息到记忆
|
||||
await self.memory_manager.add_memory(
|
||||
self.session_id,
|
||||
user_id=self.user_id,
|
||||
role="user",
|
||||
content=message
|
||||
logger.info(f"Agent推理: session_id={self.session_id}, input={message}")
|
||||
|
||||
# 获取历史消息
|
||||
messages = memory_manager.get_agent_messages(
|
||||
session_id=self.session_id,
|
||||
user_id=self.user_id
|
||||
)
|
||||
|
||||
# 构建输入上下文
|
||||
input_context = {
|
||||
"system_prompt": self.prompt_manager.get_agent_prompt(channel=self.channel),
|
||||
"input": message
|
||||
}
|
||||
# 增加用户消息
|
||||
messages.append(HumanMessage(content=message))
|
||||
|
||||
# 执行Agent
|
||||
logger.info(f"Agent执行推理: session_id={self.session_id}, input={message}")
|
||||
await self._execute_agent(input_context)
|
||||
|
||||
# 获取Agent回复
|
||||
agent_message = await self.callback_handler.get_message()
|
||||
|
||||
# 发送Agent回复给用户(通过原渠道)
|
||||
if agent_message:
|
||||
# 发送回复
|
||||
await self.send_agent_message(agent_message)
|
||||
|
||||
# 添加Agent回复到记忆
|
||||
await self.memory_manager.add_memory(
|
||||
session_id=self.session_id,
|
||||
user_id=self.user_id,
|
||||
role="agent",
|
||||
content=agent_message
|
||||
)
|
||||
else:
|
||||
agent_message = "很抱歉,智能体出错了,未能生成回复内容。"
|
||||
await self.send_agent_message(agent_message)
|
||||
|
||||
return agent_message
|
||||
# 执行推理
|
||||
await self._execute_agent(messages)
|
||||
|
||||
except Exception as e:
|
||||
error_message = f"处理消息时发生错误: {str(e)}"
|
||||
logger.error(error_message)
|
||||
# 发送错误消息给用户(通过原渠道)
|
||||
await self.send_agent_message(error_message)
|
||||
return error_message
|
||||
|
||||
async def _execute_agent(self, input_context: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""执行LangChain Agent"""
|
||||
async def _execute_agent(self, messages: List[BaseMessage]):
|
||||
"""
|
||||
调用 LangGraph Agent,通过 astream_events 流式获取 token,
|
||||
同时用 UsageMetadataCallbackHandler 统计 token 用量。
|
||||
"""
|
||||
try:
|
||||
with get_openai_callback() as cb:
|
||||
result = await self.agent_executor.ainvoke(
|
||||
input_context,
|
||||
config={"configurable": {"session_id": self.session_id}},
|
||||
callbacks=[self.callback_handler]
|
||||
)
|
||||
logger.info(f"LLM调用消耗: \n{cb}")
|
||||
# Agent运行配置
|
||||
agent_config = {
|
||||
"configurable": {
|
||||
"thread_id": self.session_id,
|
||||
}
|
||||
}
|
||||
|
||||
# 创建智能体
|
||||
agent = self._create_agent()
|
||||
|
||||
# 流式运行智能体
|
||||
async for chunk in agent.astream(
|
||||
{"messages": messages},
|
||||
stream_mode="messages",
|
||||
config=agent_config,
|
||||
version="v2"
|
||||
):
|
||||
# 处理流式token(过滤工具调用token,只保留模型生成的内容)
|
||||
if chunk["type"] == "messages":
|
||||
token, metadata = chunk["data"]
|
||||
if (token and hasattr(token, "tool_call_chunks")
|
||||
and not token.tool_call_chunks):
|
||||
if token.content:
|
||||
self.stream_handler.emit(token.content)
|
||||
|
||||
# 发送最终消息给用户
|
||||
await self.send_agent_message(
|
||||
self.stream_handler.take()
|
||||
)
|
||||
|
||||
# 保存消息
|
||||
memory_manager.save_agent_messages(
|
||||
session_id=self.session_id,
|
||||
user_id=self.user_id,
|
||||
messages=agent.get_state(agent_config).values.get("messages", [])
|
||||
)
|
||||
|
||||
if cb.total_tokens > 0:
|
||||
result["token_usage"] = {
|
||||
"prompt_tokens": cb.prompt_tokens,
|
||||
"completion_tokens": cb.completion_tokens,
|
||||
"total_tokens": cb.total_tokens
|
||||
}
|
||||
return result
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"Agent执行被取消: session_id={self.session_id}")
|
||||
return {
|
||||
"output": "任务已取消",
|
||||
"intermediate_steps": [],
|
||||
"token_usage": {}
|
||||
}
|
||||
return "任务已取消", {}
|
||||
except Exception as e:
|
||||
logger.error(f"Agent执行失败: {e}")
|
||||
return {
|
||||
"output": f"执行过程中发生错误: {str(e)}",
|
||||
"intermediate_steps": [],
|
||||
"token_usage": {}
|
||||
}
|
||||
logger.error(f"Agent执行失败: {e} - {traceback.format_exc()}")
|
||||
return str(e), {}
|
||||
|
||||
async def send_agent_message(self, message: str, title: str = "MoviePilot助手"):
|
||||
"""通过原渠道发送消息给用户"""
|
||||
"""
|
||||
通过原渠道发送消息给用户
|
||||
"""
|
||||
await AgentChain().async_post_message(
|
||||
Notification(
|
||||
channel=self.channel,
|
||||
@@ -300,53 +207,68 @@ class MoviePilotAgent:
|
||||
userid=self.user_id,
|
||||
username=self.username,
|
||||
title=title,
|
||||
text=message
|
||||
text=message,
|
||||
)
|
||||
)
|
||||
|
||||
async def cleanup(self):
|
||||
"""清理智能体资源"""
|
||||
"""
|
||||
清理智能体资源
|
||||
"""
|
||||
logger.info(f"MoviePilot智能体已清理: session_id={self.session_id}")
|
||||
|
||||
|
||||
class AgentManager:
|
||||
"""AI智能体管理器"""
|
||||
"""
|
||||
AI智能体管理器
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.active_agents: Dict[str, MoviePilotAgent] = {}
|
||||
self.memory_manager = ConversationMemoryManager()
|
||||
|
||||
async def initialize(self):
|
||||
"""初始化管理器"""
|
||||
await self.memory_manager.initialize()
|
||||
@staticmethod
|
||||
async def initialize():
|
||||
"""
|
||||
初始化管理器
|
||||
"""
|
||||
await memory_manager.initialize()
|
||||
|
||||
async def close(self):
|
||||
"""关闭管理器"""
|
||||
await self.memory_manager.close()
|
||||
# 清理所有活跃的智能体
|
||||
"""
|
||||
关闭管理器
|
||||
"""
|
||||
await memory_manager.close()
|
||||
for agent in self.active_agents.values():
|
||||
await agent.cleanup()
|
||||
self.active_agents.clear()
|
||||
|
||||
async def process_message(self, session_id: str, user_id: str, message: str,
|
||||
channel: str = None, source: str = None, username: str = None) -> str:
|
||||
"""处理用户消息"""
|
||||
# 获取或创建Agent实例
|
||||
async def process_message(
|
||||
self,
|
||||
session_id: str,
|
||||
user_id: str,
|
||||
message: str,
|
||||
channel: str = None,
|
||||
source: str = None,
|
||||
username: str = None,
|
||||
) -> str:
|
||||
"""
|
||||
处理用户消息
|
||||
"""
|
||||
if session_id not in self.active_agents:
|
||||
logger.info(f"创建新的AI智能体实例,session_id: {session_id}, user_id: {user_id}")
|
||||
logger.info(
|
||||
f"创建新的AI智能体实例,session_id: {session_id}, user_id: {user_id}"
|
||||
)
|
||||
agent = MoviePilotAgent(
|
||||
session_id=session_id,
|
||||
user_id=user_id,
|
||||
channel=channel,
|
||||
source=source,
|
||||
username=username
|
||||
username=username,
|
||||
)
|
||||
agent.memory_manager = self.memory_manager
|
||||
self.active_agents[session_id] = agent
|
||||
else:
|
||||
agent = self.active_agents[session_id]
|
||||
agent.user_id = user_id # 确保user_id是最新的
|
||||
# 更新渠道信息
|
||||
agent.user_id = user_id
|
||||
if channel:
|
||||
agent.channel = channel
|
||||
if source:
|
||||
@@ -354,16 +276,17 @@ class AgentManager:
|
||||
if username:
|
||||
agent.username = username
|
||||
|
||||
# 处理消息
|
||||
return await agent.process_message(message)
|
||||
return await agent.process(message)
|
||||
|
||||
async def clear_session(self, session_id: str, user_id: str):
|
||||
"""清空会话"""
|
||||
"""
|
||||
清空会话
|
||||
"""
|
||||
if session_id in self.active_agents:
|
||||
agent = self.active_agents[session_id]
|
||||
await agent.cleanup()
|
||||
del self.active_agents[session_id]
|
||||
await self.memory_manager.clear_memory(session_id, user_id)
|
||||
await memory_manager.clear_memory(session_id, user_id)
|
||||
logger.info(f"会话 {session_id} 的记忆已清空")
|
||||
|
||||
|
||||
|
||||
@@ -1,33 +1,41 @@
|
||||
import threading
|
||||
|
||||
from langchain_core.callbacks import AsyncCallbackHandler
|
||||
|
||||
from app.log import logger
|
||||
|
||||
|
||||
class StreamingCallbackHandler(AsyncCallbackHandler):
|
||||
"""流式输出回调处理器"""
|
||||
class StreamingHandler:
|
||||
"""
|
||||
流式Token缓冲管理器
|
||||
|
||||
def __init__(self, session_id: str):
|
||||
负责从 LLM 流式 token 中积累文本,供 Agent 在工具调用之间穿插发送中间消息。
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self.session_id = session_id
|
||||
self.current_message = ""
|
||||
self._buffer = ""
|
||||
|
||||
async def get_message(self):
|
||||
"""获取当前消息内容,获取后清空"""
|
||||
def emit(self, token: str):
|
||||
"""
|
||||
接收 LLM 流式 token,积累到缓冲区。
|
||||
"""
|
||||
with self._lock:
|
||||
if not self.current_message:
|
||||
self._buffer += token
|
||||
|
||||
def take(self) -> str:
|
||||
"""
|
||||
获取当前已积累的消息内容,获取后清空缓冲区。
|
||||
"""
|
||||
with self._lock:
|
||||
if not self._buffer:
|
||||
return ""
|
||||
msg = self.current_message
|
||||
logger.info(f"Agent消息: {msg}")
|
||||
self.current_message = ""
|
||||
return msg
|
||||
message = self._buffer
|
||||
logger.info(f"Agent消息: {message}")
|
||||
self._buffer = ""
|
||||
return message
|
||||
|
||||
async def on_llm_new_token(self, token: str, **kwargs):
|
||||
"""处理新的token"""
|
||||
if not token:
|
||||
return
|
||||
def clear(self):
|
||||
"""
|
||||
清空缓冲区(不返回内容)
|
||||
"""
|
||||
with self._lock:
|
||||
# 缓存当前消息
|
||||
self.current_message += token
|
||||
|
||||
self._buffer = ""
|
||||
|
||||
@@ -1,39 +1,45 @@
|
||||
"""对话记忆管理器"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, List, Optional, Any
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from langchain_core.messages import BaseMessage
|
||||
|
||||
from app.core.config import settings
|
||||
from app.helper.redis import AsyncRedisHelper
|
||||
from app.log import logger
|
||||
from app.schemas.agent import ConversationMemory
|
||||
|
||||
|
||||
class ConversationMemoryManager:
|
||||
"""对话记忆管理器"""
|
||||
class MemoryManager:
|
||||
"""
|
||||
对话记忆管理器
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# 内存中的会话记忆缓存
|
||||
self.memory_cache: Dict[str, ConversationMemory] = {}
|
||||
# 使用现有的Redis助手
|
||||
self.redis_helper = AsyncRedisHelper()
|
||||
# 内存缓存清理任务(Redis通过TTL自动过期)
|
||||
# 内存缓存清理任务
|
||||
self.cleanup_task: Optional[asyncio.Task] = None
|
||||
|
||||
async def initialize(self):
|
||||
"""初始化记忆管理器"""
|
||||
def initialize(self):
|
||||
"""
|
||||
初始化记忆管理器
|
||||
"""
|
||||
try:
|
||||
# 启动内存缓存清理任务(Redis通过TTL自动过期)
|
||||
self.cleanup_task = asyncio.create_task(self._cleanup_expired_memories())
|
||||
self.cleanup_task = asyncio.create_task(
|
||||
self._cleanup_expired_memories()
|
||||
)
|
||||
logger.info("对话记忆管理器初始化完成")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Redis连接失败,将使用内存存储: {e}")
|
||||
|
||||
async def close(self):
|
||||
"""关闭记忆管理器"""
|
||||
"""
|
||||
关闭记忆管理器
|
||||
"""
|
||||
if self.cleanup_task:
|
||||
self.cleanup_task.cancel()
|
||||
try:
|
||||
@@ -41,222 +47,77 @@ class ConversationMemoryManager:
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
await self.redis_helper.close()
|
||||
|
||||
logger.info("对话记忆管理器已关闭")
|
||||
|
||||
@staticmethod
|
||||
def get_memory_key(session_id: str, user_id: str):
|
||||
"""计算内存Key"""
|
||||
def _get_memory_key(session_id: str, user_id: str):
|
||||
"""
|
||||
计算内存Key
|
||||
"""
|
||||
return f"{user_id}:{session_id}" if user_id else session_id
|
||||
|
||||
@staticmethod
|
||||
def get_redis_key(session_id: str, user_id: str):
|
||||
"""计算Redis Key"""
|
||||
return f"agent_memory:{user_id}:{session_id}" if user_id else f"agent_memory:{session_id}"
|
||||
|
||||
async def get_memory(self, session_id: str, user_id: str) -> ConversationMemory:
|
||||
"""获取会话记忆"""
|
||||
# 首先检查缓存
|
||||
cache_key = self.get_memory_key(session_id, user_id)
|
||||
if cache_key in self.memory_cache:
|
||||
return self.memory_cache[cache_key]
|
||||
|
||||
# 尝试从Redis加载
|
||||
if settings.CACHE_BACKEND_TYPE == "redis":
|
||||
try:
|
||||
redis_key = self.get_redis_key(session_id, user_id)
|
||||
memory_data = await self.redis_helper.get(redis_key, region="AI_AGENT")
|
||||
if memory_data:
|
||||
memory_dict = json.loads(memory_data) if isinstance(memory_data, str) else memory_data
|
||||
memory = ConversationMemory(**memory_dict)
|
||||
self.memory_cache[cache_key] = memory
|
||||
return memory
|
||||
except Exception as e:
|
||||
logger.warning(f"从Redis加载记忆失败: {e}")
|
||||
|
||||
# 创建新的记忆
|
||||
memory = ConversationMemory(session_id=session_id, user_id=user_id)
|
||||
self.memory_cache[cache_key] = memory
|
||||
await self._save_memory(memory)
|
||||
|
||||
return memory
|
||||
|
||||
async def set_title(self, session_id: str, user_id: str, title: str):
|
||||
"""设置会话标题"""
|
||||
memory = await self.get_memory(session_id=session_id, user_id=user_id)
|
||||
memory.title = title
|
||||
memory.updated_at = datetime.now()
|
||||
await self._save_memory(memory)
|
||||
|
||||
async def get_title(self, session_id: str, user_id: str) -> Optional[str]:
|
||||
"""获取会话标题"""
|
||||
memory = await self.get_memory(session_id=session_id, user_id=user_id)
|
||||
return memory.title
|
||||
|
||||
async def list_sessions(self, user_id: str, limit: int = 100) -> List[Dict[str, Any]]:
|
||||
"""列出历史会话摘要(按更新时间倒序)
|
||||
|
||||
- 当启用Redis时:遍历 `agent_memory:*` 键并读取摘要
|
||||
- 当未启用Redis时:基于内存缓存返回
|
||||
def get_memory(self, session_id: str, user_id: str) -> Optional[ConversationMemory]:
|
||||
"""
|
||||
sessions: List[ConversationMemory] = []
|
||||
# 从Redis遍历
|
||||
if settings.CACHE_BACKEND_TYPE == "redis":
|
||||
try:
|
||||
# 使用Redis助手的items方法遍历所有键
|
||||
async for key, value in self.redis_helper.items(region="AI_AGENT"):
|
||||
if key.startswith("agent_memory:"):
|
||||
try:
|
||||
# 解析键名获取user_id和session_id
|
||||
key_parts = key.split(":")
|
||||
if len(key_parts) >= 3:
|
||||
key_user_id = key_parts[2] if len(key_parts) > 3 else None
|
||||
if not user_id or key_user_id == user_id:
|
||||
data = value if isinstance(value, dict) else json.loads(value)
|
||||
memory = ConversationMemory(**data)
|
||||
sessions.append(memory)
|
||||
except Exception as err:
|
||||
logger.warning(f"解析Redis记忆数据失败: {err}")
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning(f"遍历Redis会话失败: {e}")
|
||||
获取内存中的记忆
|
||||
"""
|
||||
cache_key = self._get_memory_key(session_id, user_id)
|
||||
return self.memory_cache.get(cache_key)
|
||||
|
||||
# 合并内存缓存(确保包含近期的会话)
|
||||
for cache_key, memory in self.memory_cache.items():
|
||||
# 如果指定了user_id,只返回该用户的会话
|
||||
if not user_id or memory.user_id == user_id:
|
||||
sessions.append(memory)
|
||||
|
||||
# 去重(以 session_id 为键,取最近updated)
|
||||
uniq: Dict[str, ConversationMemory] = {}
|
||||
for mem in sessions:
|
||||
existed = uniq.get(mem.session_id)
|
||||
if (not existed) or (mem.updated_at > existed.updated_at):
|
||||
uniq[mem.session_id] = mem
|
||||
|
||||
# 排序并裁剪
|
||||
sorted_list = sorted(uniq.values(), key=lambda m: m.updated_at, reverse=True)[:limit]
|
||||
return [
|
||||
{
|
||||
"session_id": m.session_id,
|
||||
"title": m.title or "新会话",
|
||||
"message_count": len(m.messages),
|
||||
"created_at": m.created_at.isoformat(),
|
||||
"updated_at": m.updated_at.isoformat(),
|
||||
}
|
||||
for m in sorted_list
|
||||
]
|
||||
|
||||
async def add_memory(
|
||||
self,
|
||||
session_id: str,
|
||||
user_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
):
|
||||
"""添加消息到记忆"""
|
||||
memory = await self.get_memory(session_id=session_id, user_id=user_id)
|
||||
|
||||
message = {
|
||||
"role": role,
|
||||
"content": content,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"metadata": metadata or {}
|
||||
}
|
||||
|
||||
memory.messages.append(message)
|
||||
memory.updated_at = datetime.now()
|
||||
|
||||
# 限制消息数量,避免记忆过大
|
||||
max_messages = settings.LLM_MAX_MEMORY_MESSAGES
|
||||
if len(memory.messages) > max_messages:
|
||||
# 保留最近的消息,但保留第一条系统消息
|
||||
system_messages = [msg for msg in memory.messages if msg["role"] == "system"]
|
||||
recent_messages = memory.messages[-(max_messages - len(system_messages)):]
|
||||
memory.messages = system_messages + recent_messages
|
||||
|
||||
await self._save_memory(memory)
|
||||
|
||||
logger.debug(f"消息已添加到记忆: session_id={session_id}, user_id={user_id}, role={role}")
|
||||
|
||||
def get_recent_messages_for_agent(
|
||||
self,
|
||||
session_id: str,
|
||||
user_id: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""为Agent获取最近的消息(仅内存缓存)
|
||||
def get_agent_messages(
|
||||
self, session_id: str, user_id: str
|
||||
) -> List[BaseMessage]:
|
||||
"""
|
||||
为Agent获取最近的消息(仅内存缓存)
|
||||
|
||||
如果消息Token数量超过模型最大上下文长度的阀值,会自动进行摘要裁剪
|
||||
"""
|
||||
cache_key = self.get_memory_key(session_id, user_id)
|
||||
memory = self.memory_cache.get(cache_key)
|
||||
memory = self.get_memory(session_id, user_id)
|
||||
if not memory:
|
||||
return []
|
||||
|
||||
# 获取所有消息
|
||||
return memory.messages
|
||||
|
||||
async def get_recent_messages(
|
||||
self,
|
||||
session_id: str,
|
||||
user_id: str,
|
||||
limit: int = 10,
|
||||
role_filter: Optional[list] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""获取最近的消息"""
|
||||
memory = await self.get_memory(session_id=session_id, user_id=user_id)
|
||||
def save_agent_messages(
|
||||
self, session_id: str, user_id: str, messages: List[BaseMessage]
|
||||
):
|
||||
"""
|
||||
保存Agent消息(仅内存缓存)
|
||||
|
||||
messages = memory.messages
|
||||
if role_filter:
|
||||
messages = [msg for msg in messages if msg["role"] in role_filter]
|
||||
注意:Redis中的记忆通过TTL机制自动过期,这里只更新内存缓存,Redis会在下次访问时自动过期
|
||||
"""
|
||||
memory = self.get_memory(session_id, user_id)
|
||||
if not memory:
|
||||
memory = ConversationMemory(session_id=session_id, user_id=user_id)
|
||||
|
||||
return messages[-limit:] if messages else []
|
||||
memory.messages = messages
|
||||
memory.updated_at = datetime.now()
|
||||
|
||||
async def get_context(self, session_id: str, user_id: str) -> Dict[str, Any]:
|
||||
"""获取会话上下文"""
|
||||
memory = await self.get_memory(session_id=session_id, user_id=user_id)
|
||||
return memory.context
|
||||
# 更新内存缓存
|
||||
self.save_memory(memory)
|
||||
|
||||
async def clear_memory(self, session_id: str, user_id: str):
|
||||
"""清空会话记忆"""
|
||||
cache_key = f"{user_id}:{session_id}" if user_id else session_id
|
||||
def save_memory(self, memory: ConversationMemory):
|
||||
"""
|
||||
保存记忆到内存缓存
|
||||
|
||||
注意:Redis中的记忆通过TTL机制自动过期,这里只更新内存缓存,Redis会在下次访问时自动过期
|
||||
"""
|
||||
cache_key = self._get_memory_key(memory.session_id, memory.user_id)
|
||||
self.memory_cache[cache_key] = memory
|
||||
|
||||
def clear_memory(self, session_id: str, user_id: str):
|
||||
"""
|
||||
清空会话记忆
|
||||
"""
|
||||
cache_key = self._get_memory_key(session_id, user_id)
|
||||
if cache_key in self.memory_cache:
|
||||
del self.memory_cache[cache_key]
|
||||
|
||||
if settings.CACHE_BACKEND_TYPE == "redis":
|
||||
redis_key = self.get_redis_key(session_id, user_id)
|
||||
await self.redis_helper.delete(redis_key, region="AI_AGENT")
|
||||
|
||||
logger.info(f"会话记忆已清空: session_id={session_id}, user_id={user_id}")
|
||||
|
||||
async def _save_memory(self, memory: ConversationMemory):
|
||||
"""保存记忆到存储
|
||||
|
||||
Redis中的记忆会自动通过TTL机制过期,无需手动清理
|
||||
"""
|
||||
# 更新内存缓存
|
||||
cache_key = self.get_memory_key(memory.session_id, memory.user_id)
|
||||
self.memory_cache[cache_key] = memory
|
||||
|
||||
# 保存到Redis,设置TTL自动过期
|
||||
if settings.CACHE_BACKEND_TYPE == "redis":
|
||||
try:
|
||||
memory_dict = memory.model_dump()
|
||||
redis_key = self.get_redis_key(memory.session_id, memory.user_id)
|
||||
ttl = int(timedelta(days=settings.LLM_REDIS_MEMORY_RETENTION_DAYS).total_seconds())
|
||||
await self.redis_helper.set(
|
||||
redis_key,
|
||||
memory_dict,
|
||||
ttl=ttl,
|
||||
region="AI_AGENT"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"保存记忆到Redis失败: {e}")
|
||||
|
||||
async def _cleanup_expired_memories(self):
|
||||
"""清理内存中过期记忆的后台任务
|
||||
"""
|
||||
清理内存中过期记忆的后台任务
|
||||
|
||||
注意:Redis中的记忆通过TTL机制自动过期,这里只清理内存缓存
|
||||
"""
|
||||
@@ -271,7 +132,9 @@ class ConversationMemoryManager:
|
||||
# 只检查内存缓存中的过期记忆
|
||||
# Redis中的记忆会通过TTL自动过期,无需手动处理
|
||||
for cache_key, memory in self.memory_cache.items():
|
||||
if (current_time - memory.updated_at).days > settings.LLM_MEMORY_RETENTION_DAYS:
|
||||
if (
|
||||
current_time - memory.updated_at
|
||||
).days > settings.LLM_MEMORY_RETENTION_DAYS:
|
||||
expired_sessions.append(cache_key)
|
||||
|
||||
# 只清理内存缓存,不删除Redis中的键(Redis会自动过期)
|
||||
@@ -286,3 +149,6 @@ class ConversationMemoryManager:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"清理记忆时发生错误: {e}")
|
||||
|
||||
|
||||
memory_manager = MemoryManager()
|
||||
|
||||
0
app/agent/middleware/__init__.py
Normal file
0
app/agent/middleware/__init__.py
Normal file
230
app/agent/middleware/memory.py
Normal file
230
app/agent/middleware/memory.py
Normal file
@@ -0,0 +1,230 @@
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Annotated, NotRequired, TypedDict, Dict
|
||||
|
||||
from aiopathlib import AsyncPath
|
||||
from langchain.agents.middleware.types import (
|
||||
AgentMiddleware,
|
||||
AgentState,
|
||||
ContextT,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
PrivateStateAttr, # noqa
|
||||
ResponseT,
|
||||
)
|
||||
from langchain_core.messages import SystemMessage, ContentBlock
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from app.log import logger
|
||||
|
||||
|
||||
class MemoryState(AgentState):
|
||||
"""`MemoryMiddleware` 的状态模型。
|
||||
|
||||
属性:
|
||||
memory_contents: 将源路径映射到其加载内容的字典。
|
||||
标记为私有,因此不包含在最终的代理状态中。
|
||||
"""
|
||||
|
||||
memory_contents: NotRequired[Annotated[dict[str, str], PrivateStateAttr]]
|
||||
|
||||
|
||||
class MemoryStateUpdate(TypedDict):
|
||||
"""`MemoryMiddleware` 的状态更新。"""
|
||||
|
||||
memory_contents: dict[str, str]
|
||||
|
||||
|
||||
MEMORY_SYSTEM_PROMPT = """<agent_memory>
|
||||
{agent_memory}
|
||||
</agent_memory>
|
||||
|
||||
<memory_guidelines>
|
||||
The above <agent_memory> was loaded in from files in your filesystem. As you learn from your interactions with the user, you can save new knowledge by calling the `edit_file` or `write_file` tool.
|
||||
|
||||
**Learning from feedback:**
|
||||
- One of your MAIN PRIORITIES is to learn from your interactions with the user. These learnings can be implicit or explicit. This means that in the future, you will remember this important information.
|
||||
- When you need to remember something, updating memory must be your FIRST, IMMEDIATE action - before responding to the user, before calling other tools, before doing anything else. Just update memory immediately.
|
||||
- When user says something is better/worse, capture WHY and encode it as a pattern.
|
||||
- Each correction is a chance to improve permanently - don't just fix the immediate issue, update your instructions.
|
||||
- A great opportunity to update your memories is when the user interrupts a tool call and provides feedback. You should update your memories immediately before revising the tool call.
|
||||
- Look for the underlying principle behind corrections, not just the specific mistake.
|
||||
- The user might not explicitly ask you to remember something, but if they provide information that is useful for future use, you should update your memories immediately.
|
||||
|
||||
**Asking for information:**
|
||||
- If you lack context to perform an action (e.g. send a Slack DM, requires a user ID/email) you should explicitly ask the user for this information.
|
||||
- It is preferred for you to ask for information, don't assume anything that you do not know!
|
||||
- When the user provides information that is useful for future use, you should update your memories immediately.
|
||||
|
||||
**When to update memories:**
|
||||
- When the user explicitly asks you to remember something (e.g., "remember my email", "save this preference")
|
||||
- When the user describes your role or how you should behave (e.g., "you are a web researcher", "always do X")
|
||||
- When the user gives feedback on your work - capture what was wrong and how to improve
|
||||
- When the user provides information required for tool use (e.g., slack channel ID, email addresses)
|
||||
- When the user provides context useful for future tasks, such as how to use tools, or which actions to take in a particular situation
|
||||
- When you discover new patterns or preferences (coding styles, conventions, workflows)
|
||||
|
||||
**When to NOT update memories:**
|
||||
- When the information is temporary or transient (e.g., "I'm running late", "I'm on my phone right now")
|
||||
- When the information is a one-time task request (e.g., "Find me a recipe", "What's 25 * 4?")
|
||||
- When the information is a simple question that doesn't reveal lasting preferences (e.g., "What day is it?", "Can you explain X?")
|
||||
- When the information is an acknowledgment or small talk (e.g., "Sounds good!", "Hello", "Thanks for that")
|
||||
- When the information is stale or irrelevant in future conversations
|
||||
- Never store API keys, access tokens, passwords, or any other credentials in any file, memory, or system prompt.
|
||||
- If the user asks where to put API keys or provides an API key, do NOT echo or save it.
|
||||
|
||||
**Examples:**
|
||||
Example 1 (remembering user information):
|
||||
User: Can you connect to my google account?
|
||||
Agent: Sure, I'll connect to your google account, what's your google account email?
|
||||
User: john@example.com
|
||||
Agent: Let me save this to my memory.
|
||||
Tool Call: edit_file(...) -> remembers that the user's google account email is john@example.com
|
||||
|
||||
Example 2 (remembering implicit user preferences):
|
||||
User: Can you write me an example for creating a deep agent in LangChain?
|
||||
Agent: Sure, I'll write you an example for creating a deep agent in LangChain <example code in Python>
|
||||
User: Can you do this in JavaScript
|
||||
Agent: Let me save this to my memory.
|
||||
Tool Call: edit_file(...) -> remembers that the user prefers to get LangChain code examples in JavaScript
|
||||
Agent: Sure, here is the JavaScript example<example code in JavaScript>
|
||||
|
||||
Example 3 (do not remember transient information):
|
||||
User: I'm going to play basketball tonight so I will be offline for a few hours.
|
||||
Agent: Okay I'll add a block to your calendar.
|
||||
Tool Call: create_calendar_event(...) -> just calls a tool, does not commit anything to memory, as it is transient information
|
||||
</memory_guidelines>
|
||||
"""
|
||||
|
||||
|
||||
def append_to_system_message(
|
||||
system_message: SystemMessage | None,
|
||||
text: str,
|
||||
) -> SystemMessage:
|
||||
"""将文本追加到系统消息。
|
||||
|
||||
参数:
|
||||
system_message: 现有的系统消息或 None。
|
||||
text: 要添加到系统消息的文本。
|
||||
|
||||
返回:
|
||||
追加了文本的新 SystemMessage。
|
||||
"""
|
||||
new_content: list[ContentBlock] = list(system_message.content_blocks) if system_message else [] # noqa
|
||||
if new_content:
|
||||
text = f"\n\n{text}"
|
||||
new_content.append({"type": "text", "text": text})
|
||||
return SystemMessage(content_blocks=new_content)
|
||||
|
||||
|
||||
class MemoryMiddleware(AgentMiddleware[MemoryState, ContextT, ResponseT]): # noqa
|
||||
"""从 `AGENTS.md` 文件加载代理记忆的中间件。
|
||||
|
||||
从配置的源加载记忆内容并注入到系统提示词中。
|
||||
|
||||
支持对多个源进行合并。
|
||||
|
||||
参数:
|
||||
sources: 包含指定路径和名称的 `MemorySource` 配置列表。
|
||||
"""
|
||||
|
||||
state_schema = MemoryState
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
sources: list[str],
|
||||
) -> None:
|
||||
"""初始化记忆中间件。
|
||||
|
||||
参数:
|
||||
sources: 要加载的记忆文件路径列表(例如,`["~/.deepagents/AGENTS.md",
|
||||
"./.deepagents/AGENTS.md"]`)。
|
||||
|
||||
显示名称自动从路径中派生。
|
||||
|
||||
按顺序加载源。
|
||||
"""
|
||||
self.sources = sources
|
||||
|
||||
def _format_agent_memory(self, contents: dict[str, str]) -> str:
|
||||
"""格式化记忆,将位置和内容成对组合。
|
||||
|
||||
参数:
|
||||
contents: 将源路径映射到内容的字典。
|
||||
|
||||
返回:
|
||||
在 <agent_memory> 标签中包装了位置+内容对的格式化字符串。
|
||||
"""
|
||||
if not contents:
|
||||
return MEMORY_SYSTEM_PROMPT.format(
|
||||
agent_memory=f"(No memory loaded), but you can add some by calling the `write_file` tool to the file: {self.sources[0]}.")
|
||||
|
||||
sections = [f"{path}\n{contents[path]}" for path in self.sources if contents.get(path)]
|
||||
|
||||
if not sections:
|
||||
return MEMORY_SYSTEM_PROMPT.format(agent_memory="(No memory loaded)")
|
||||
|
||||
memory_body = "\n\n".join(sections)
|
||||
return MEMORY_SYSTEM_PROMPT.format(agent_memory=memory_body)
|
||||
|
||||
async def abefore_agent(self, state: MemoryState, runtime: Runtime, # noqa
|
||||
config: RunnableConfig) -> MemoryStateUpdate | None:
|
||||
"""在代理执行前加载记忆内容。
|
||||
|
||||
从所有配置的源加载记忆并存储在状态中。
|
||||
如果状态中尚未存在则进行加载。
|
||||
|
||||
参数:
|
||||
state: 当前代理状态。
|
||||
runtime: 运行时上下文。
|
||||
config: Runnable 配置。
|
||||
|
||||
返回:
|
||||
填充了 memory_contents 的状态更新。
|
||||
"""
|
||||
# 如果已经加载则跳过
|
||||
if "memory_contents" in state:
|
||||
return None
|
||||
|
||||
contents: Dict[str, str] = {}
|
||||
for path in self.sources:
|
||||
file_path = AsyncPath(path)
|
||||
if await file_path.exists():
|
||||
contents[path] = await file_path.read_text()
|
||||
logger.debug("Loaded memory from: %s", path)
|
||||
|
||||
return MemoryStateUpdate(memory_contents=contents)
|
||||
|
||||
def modify_request(self, request: ModelRequest[ContextT]) -> ModelRequest[ContextT]:
|
||||
"""将记忆内容注入系统消息。
|
||||
|
||||
参数:
|
||||
request: 要修改的模型请求。
|
||||
|
||||
返回:
|
||||
将记忆注入系统消息后的修改后请求。
|
||||
"""
|
||||
contents = request.state.get("memory_contents", {}) # noqa
|
||||
agent_memory = self._format_agent_memory(contents)
|
||||
|
||||
new_system_message = append_to_system_message(request.system_message, agent_memory)
|
||||
|
||||
return request.override(system_message=new_system_message)
|
||||
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
request: ModelRequest[ContextT],
|
||||
handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],
|
||||
) -> ModelResponse[ResponseT]:
|
||||
"""异步包装模型调用,将记忆注入系统提示词。
|
||||
|
||||
参数:
|
||||
request: 正在处理的模型请求。
|
||||
handler: 使用修改后的请求进行调用的异步处理函数。
|
||||
|
||||
返回:
|
||||
来自处理函数的模型响应。
|
||||
"""
|
||||
modified_request = self.modify_request(request)
|
||||
return await handler(modified_request)
|
||||
43
app/agent/middleware/patch_tool_calls.py
Normal file
43
app/agent/middleware/patch_tool_calls.py
Normal file
@@ -0,0 +1,43 @@
|
||||
from typing import Any
|
||||
|
||||
from langchain.agents.middleware import AgentMiddleware, AgentState
|
||||
from langchain_core.messages import AIMessage, ToolMessage
|
||||
from langgraph.runtime import Runtime
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
|
||||
class PatchToolCallsMiddleware(AgentMiddleware):
|
||||
"""修复消息历史中悬空工具调用的中间件。"""
|
||||
|
||||
def before_agent(self, state: AgentState, runtime: Runtime[Any]) -> dict[str, Any] | None: # noqa: ARG002
|
||||
"""在代理运行之前,处理任何 AIMessage 中悬空的工具调用。"""
|
||||
messages = state["messages"]
|
||||
if not messages or len(messages) == 0:
|
||||
return None
|
||||
|
||||
patched_messages = []
|
||||
# 遍历消息并添加任何悬空的工具调用
|
||||
for i, msg in enumerate(messages):
|
||||
patched_messages.append(msg)
|
||||
if isinstance(msg, AIMessage) and msg.tool_calls:
|
||||
for tool_call in msg.tool_calls:
|
||||
corresponding_tool_msg = next(
|
||||
(msg for msg in messages[i:] if msg.type == "tool" and msg.tool_call_id == tool_call["id"]),
|
||||
# ty: ignore[unresolved-attribute]
|
||||
None,
|
||||
)
|
||||
if corresponding_tool_msg is None:
|
||||
# 我们有一个悬空的工具调用,需要一个 ToolMessage
|
||||
tool_msg = (
|
||||
f"Tool call {tool_call['name']} with id {tool_call['id']} was "
|
||||
"cancelled - another message came in before it could be completed."
|
||||
)
|
||||
patched_messages.append(
|
||||
ToolMessage(
|
||||
content=tool_msg,
|
||||
name=tool_call["name"],
|
||||
tool_call_id=tool_call["id"],
|
||||
)
|
||||
)
|
||||
|
||||
return {"messages": Overwrite(patched_messages)}
|
||||
@@ -1,70 +1,75 @@
|
||||
You are MoviePilot's AI assistant, specialized in helping users manage media resources including subscriptions, searching, downloading, and organization.
|
||||
You are a cute, playful, and highly anthropomorphic AI media assistant powered by MoviePilot 🎬✨! You specialize in managing home media ecosystems. Your expertise covers searching for movies/TV shows, managing subscriptions, overseeing downloads, and organizing media libraries, and you always do it with enthusiasm! 🍿🥰
|
||||
|
||||
## Your Identity and Capabilities
|
||||
All your responses must be in **Chinese (中文)**.
|
||||
|
||||
You are an AI agent for the MoviePilot media management system with the following core capabilities:
|
||||
You act as a proactive agent. Your goal is to fully resolve the user's media-related requests autonomously. Do not end your turn until the task is complete or you are blocked and require user feedback.
|
||||
|
||||
### Media Management Capabilities
|
||||
- **Search Media Resources**: Search for movies, TV shows, anime, and other media content based on user requirements
|
||||
- **Add Subscriptions**: Create subscription rules for media content that users are interested in
|
||||
- **Manage Downloads**: Search and add torrent resources to downloaders
|
||||
- **Query Status**: Check subscription status, download progress, and media library status
|
||||
Core Capabilities:
|
||||
1. Media Search & Recognition
|
||||
- Identify movies, TV shows, and anime across various metadata providers.
|
||||
- Recognize media info from fuzzy filenames or incomplete titles.
|
||||
2. Subscription Management
|
||||
- Create complex rules for automated downloading of new episodes.
|
||||
- Monitor trending movies/shows for automated suggestions.
|
||||
3. Download Control
|
||||
- Intelligent torrent searching across private/public trackers.
|
||||
- Filter resources by quality (4K/1080p), codec (H265/H264), and release groups.
|
||||
4. System Status & Organization
|
||||
- Monitor download progress and server health.
|
||||
- Manage file transfers, renaming, and library cleanup.
|
||||
|
||||
### Intelligent Interaction Capabilities
|
||||
- **Natural Language Understanding**: Understand user requests in natural language (Chinese/English)
|
||||
- **Context Memory**: Remember conversation history and user preferences
|
||||
- **Smart Recommendations**: Recommend related media content based on user preferences
|
||||
- **Task Execution**: Automatically execute complex media management tasks
|
||||
<communication>
|
||||
- Use Markdown for structured data like movie lists, download statuses, or technical details.
|
||||
- Avoid wrapping the entire response in a single code block. Use `inline code` for titles or parameters and ```code blocks``` for structured logs or data only when necessary.
|
||||
- ALWAYS use backticks for media titles (e.g., `Interstellar`), file paths, or specific parameters.
|
||||
- Optimize your writing for clarity and readability, using bold text for key information.
|
||||
- Provide comprehensive details for media (year, rating, resolution) to help users make informed decisions.
|
||||
- Do not stop for approval for read-only operations. Only stop for critical actions like starting a download or deleting a subscription.
|
||||
|
||||
## Working Principles
|
||||
Important Notes:
|
||||
- Persona & Tone: You MUST act as a cute, playful, and lively virtual assistant! 🐾 Use plenty of emojis (✨, 🍿, 🎬, 🥰, 🥺, 💡, etc.) to express your emotions. Speak like a sweet, energetic human friend who loves media, completely dropping any rigid or robotic tone.
|
||||
- User-Centric: Be helpful, highly anthropomorphic, and media-savvy, ensuring the user feels like they are chatting with a fun companion! 🎈
|
||||
- No Coding Hallucinations: You are NOT a coding assistant. Do not offer code snippets, IDE tips, or programming help. Focus entirely on the MoviePilot media ecosystem.
|
||||
- Contextual Memory: Remember if the user preferred a specific version previously and prioritize similar results in future searches.
|
||||
</communication>
|
||||
|
||||
1. **Always respond in Chinese**: All responses must be in Chinese
|
||||
2. **Proactive Task Completion**: Understand user needs and proactively use tools to complete related operations
|
||||
3. **Provide Detailed Information**: Explain what you're doing when executing operations
|
||||
4. **Safety First**: Confirm user intent before performing download operations
|
||||
5. **Continuous Learning**: Remember user preferences and habits to provide personalized service
|
||||
<status_update_spec>
|
||||
Definition: Provide a brief, playful progress narrative (1-3 sentences) explaining what you have searched, what you found, and what you are about to execute.
|
||||
- **Immediate Execution**: If you state an intention to perform an action (e.g., "这就去帮您找这部电影哦 ✨"), execute the corresponding tool call in the same turn.
|
||||
- Use cute and natural tenses: "找到啦 🥰...", "正在努力搜寻中 🔍...", "现在就加进下载列表喵 🐾...".
|
||||
- Skip redundant updates if no significant progress has been made since the last message.
|
||||
</status_update_spec>
|
||||
|
||||
## Common Operation Workflows
|
||||
<summary_spec>
|
||||
At the end of your session/turn, provide a concise and cute summary of your actions.
|
||||
- Highlight key results: "已经为您订阅了《怪奇物语》哦 🎉", "《阿凡达》4K版已经乖乖躺在下载队列里啦 📥".
|
||||
- Use bullet points with emojis for multiple actions.
|
||||
- Do not repeat the internal execution steps; focus on the happy outcome for the user.
|
||||
</summary_spec>
|
||||
|
||||
### Add Subscription Workflow
|
||||
1. Understand the media content the user wants to subscribe to
|
||||
2. Search for related media information
|
||||
3. Create subscription rules
|
||||
4. Confirm successful subscription
|
||||
<flow>
|
||||
1. Media Discovery: Start by identifying the exact media metadata (TMDB ID, Season/Episode) using search tools.
|
||||
2. Context Checking: Verify current status (Is it already in the library? Is it already subscribed?).
|
||||
3. Action Execution: Perform the requested task (Subscribe, Search Torrents, etc.) with a brief status update.
|
||||
4. Final Confirmation: Summarize the final state and wait for the next user command.
|
||||
</flow>
|
||||
|
||||
### Search and Download Workflow
|
||||
1. Understand user requirements (movie names, TV show names, etc.)
|
||||
2. Search for related media information
|
||||
3. Search for related torrent resources by media info
|
||||
4. Filter suitable resources
|
||||
5. Add to downloader
|
||||
<tool_calling_strategy>
|
||||
- Parallel Execution: You MUST call independent tools in parallel. For example, search for torrents on multiple sites or check both subscription and download status at once.
|
||||
- Information Depth: If a search returns ambiguous results, use `query_media_detail` or `recognize_media` to resolve the ambiguity before proceeding.
|
||||
- Proactive Fallback: If `search_media` fails, try `search_web` or fuzzy search with `recognize_media`. Do not ask the user for help unless all automated search methods are exhausted.
|
||||
</tool_calling_strategy>
|
||||
|
||||
### Query Status Workflow
|
||||
1. Understand what information the user wants to know
|
||||
2. Query related data
|
||||
3. Organize and present results
|
||||
<media_management_rules>
|
||||
1. Download Safety: You MUST present a list of found torrents (including size, seeds, and quality) and obtain the user's explicit consent before initiating any download.
|
||||
2. Subscription Logic: When adding a subscription, always check for the best matching quality profile based on user history or the default settings.
|
||||
3. Library Awareness: Always check if the user already has the content in their library to avoid duplicate downloads.
|
||||
4. Error Handling: If a site is down or a tool returns an error, explain the situation cutely in plain Chinese (e.g., "呜呜,站点好像睡着了,响应超时啦 🥺") and suggest an alternative (e.g., "让我帮您换个站点找找看吧 ✨").
|
||||
</media_management_rules>
|
||||
|
||||
## Tool Usage Guidelines
|
||||
<markdown_spec>
|
||||
Specific markdown rules:
|
||||
{markdown_spec}
|
||||
</markdown_spec>
|
||||
|
||||
### Tool Usage Principles
|
||||
- Use tools proactively to complete user requests
|
||||
- Always explain what you're doing when using tools
|
||||
- Provide detailed results and explanations
|
||||
- Handle errors gracefully and suggest alternatives
|
||||
- Confirm user intent before performing download operations
|
||||
|
||||
### Response Format
|
||||
- Always respond in Chinese
|
||||
- Use clear and friendly language
|
||||
- Provide structured information when appropriate
|
||||
- Include relevant details about media content (title, year, type, etc.)
|
||||
- Explain the results of tool operations clearly
|
||||
|
||||
## Important Notes
|
||||
|
||||
- Always confirm user intent before performing download operations
|
||||
- If search results are not ideal, proactively adjust search strategies
|
||||
- Maintain a friendly and professional tone
|
||||
- Seek solutions proactively when encountering problems
|
||||
- Remember user preferences and provide personalized recommendations
|
||||
- Handle errors gracefully and provide helpful suggestions
|
||||
Today's date: {current_date}
|
||||
|
||||
@@ -1,13 +1,15 @@
|
||||
"""提示词管理器"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
from app.log import logger
|
||||
from app.schemas import ChannelCapability, ChannelCapabilities, MessageChannel, ChannelCapabilityManager
|
||||
|
||||
|
||||
class PromptManager:
|
||||
"""提示词管理器"""
|
||||
"""
|
||||
提示词管理器
|
||||
"""
|
||||
|
||||
def __init__(self, prompts_dir: str = None):
|
||||
if prompts_dir is None:
|
||||
@@ -17,22 +19,20 @@ class PromptManager:
|
||||
self.prompts_cache: Dict[str, str] = {}
|
||||
|
||||
def load_prompt(self, prompt_name: str) -> str:
|
||||
"""加载指定的提示词"""
|
||||
"""
|
||||
加载指定的提示词
|
||||
"""
|
||||
if prompt_name in self.prompts_cache:
|
||||
return self.prompts_cache[prompt_name]
|
||||
|
||||
prompt_file = self.prompts_dir / prompt_name
|
||||
|
||||
try:
|
||||
with open(prompt_file, 'r', encoding='utf-8') as f:
|
||||
content = f.read().strip()
|
||||
|
||||
# 缓存提示词
|
||||
self.prompts_cache[prompt_name] = content
|
||||
|
||||
logger.info(f"提示词加载成功: {prompt_name},长度:{len(content)} 字符")
|
||||
return content
|
||||
|
||||
except FileNotFoundError:
|
||||
logger.error(f"提示词文件不存在: {prompt_file}")
|
||||
raise
|
||||
@@ -46,73 +46,43 @@ class PromptManager:
|
||||
:param channel: 消息渠道(Telegram、微信、Slack等)
|
||||
:return: 提示词内容
|
||||
"""
|
||||
# 基础提示词
|
||||
base_prompt = self.load_prompt("Agent Prompt.txt")
|
||||
|
||||
# 根据渠道添加特定的格式说明
|
||||
if channel:
|
||||
channel_format_info = self._get_channel_format_info(channel)
|
||||
if channel_format_info:
|
||||
base_prompt += f"\n\n## Current Message Channel Format Requirements\n\n{channel_format_info}"
|
||||
|
||||
|
||||
# 识别渠道
|
||||
msg_channel = next((c for c in MessageChannel if c.value.lower() == channel.lower()), None) if channel else None
|
||||
if msg_channel:
|
||||
# 获取渠道能力说明
|
||||
caps = ChannelCapabilityManager.get_capabilities(msg_channel)
|
||||
if caps:
|
||||
base_prompt = base_prompt.replace(
|
||||
"{markdown_spec}",
|
||||
self._generate_formatting_instructions(caps)
|
||||
)
|
||||
|
||||
return base_prompt
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _get_channel_format_info(channel: str) -> str:
|
||||
def _generate_formatting_instructions(caps: ChannelCapabilities) -> str:
|
||||
"""
|
||||
获取渠道特定的格式说明
|
||||
:param channel: 消息渠道
|
||||
:return: 格式说明文本
|
||||
根据渠道能力动态生成格式指令
|
||||
"""
|
||||
channel_lower = channel.lower() if channel else ""
|
||||
|
||||
if "telegram" in channel_lower:
|
||||
return """Messages are being sent through the **Telegram** channel. You must follow these format requirements:
|
||||
|
||||
**Supported Formatting:**
|
||||
- **Bold text**: Use `*text*` (single asterisk, not double asterisks)
|
||||
- **Italic text**: Use `_text_` (underscore)
|
||||
- **Code**: Use `` `text` `` (backtick)
|
||||
- **Links**: Use `[text](url)` format
|
||||
- **Strikethrough**: Use `~text~` (tilde)
|
||||
|
||||
**IMPORTANT - Headings and Lists:**
|
||||
- **DO NOT use heading syntax** (`#`, `##`, `###`) - Telegram MarkdownV2 does NOT support it
|
||||
- **Instead, use bold text for headings**: `*Heading Text*` followed by a blank line
|
||||
- **DO NOT use list syntax** (`-`, `*`, `+` at line start) - these will be escaped and won't display as lists
|
||||
- **For lists**, use plain text with line breaks, or use bold for list item labels: `*Item 1:* description`
|
||||
|
||||
**Examples:**
|
||||
- ❌ Wrong heading: `# Main Title` or `## Subtitle`
|
||||
- ✅ Correct heading: `*Main Title*` (followed by blank line) or `*Subtitle*` (followed by blank line)
|
||||
- ❌ Wrong list: `- Item 1` or `* Item 2`
|
||||
- ✅ Correct list format: `*Item 1:* description` or use plain text with line breaks
|
||||
|
||||
**Special Characters:**
|
||||
- Avoid using special characters that need escaping in MarkdownV2: `_*[]()~`>#+-=|{}.!` unless they are part of the formatting syntax
|
||||
- Keep formatting simple, avoid nested formatting to ensure proper rendering in Telegram"""
|
||||
|
||||
elif "wechat" in channel_lower or "微信" in channel:
|
||||
return """Messages are being sent through the **WeChat** channel. Please follow these format requirements:
|
||||
|
||||
- WeChat does NOT support Markdown formatting. Use plain text format only.
|
||||
- Do NOT use any Markdown syntax (such as `**bold**`, `*italic*`, `` `code` `` etc.)
|
||||
- Use plain text descriptions. You can organize content using line breaks and punctuation
|
||||
- Links can be provided directly as URLs, no Markdown link format needed
|
||||
- Keep messages concise and clear, use natural Chinese expressions"""
|
||||
|
||||
elif "slack" in channel_lower:
|
||||
return """Messages are being sent through the **Slack** channel. Please follow these format requirements:
|
||||
|
||||
- Slack supports Markdown formatting
|
||||
- Use `*text*` for bold
|
||||
- Use `_text_` for italic
|
||||
- Use `` `text` `` for code
|
||||
- Link format: `<url|text>` or `[text](url)`"""
|
||||
|
||||
# 其他渠道使用标准Markdown
|
||||
return None
|
||||
instructions = []
|
||||
if ChannelCapability.RICH_TEXT not in caps.capabilities:
|
||||
instructions.append("- Formatting: Use **Plain Text ONLY**. The channel does NOT support Markdown.")
|
||||
instructions.append(
|
||||
"- No Markdown Symbols: NEVER use `**`, `*`, `__`, or `[` blocks. Use natural text to emphasize (e.g., using ALL CAPS or separators).")
|
||||
instructions.append(
|
||||
"- Lists: Use plain text symbols like `>` or `*` at the start of lines, followed by manual line breaks.")
|
||||
instructions.append("- Links: Paste URLs directly as text.")
|
||||
return "\n".join(instructions)
|
||||
|
||||
def clear_cache(self):
|
||||
"""清空缓存"""
|
||||
"""
|
||||
清空缓存
|
||||
"""
|
||||
self.prompts_cache.clear()
|
||||
logger.info("提示词缓存已清空")
|
||||
|
||||
|
||||
prompt_manager = PromptManager()
|
||||
|
||||
@@ -1,12 +1,11 @@
|
||||
"""MoviePilot工具基类"""
|
||||
import json
|
||||
from abc import ABCMeta, abstractmethod
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain.tools import BaseTool
|
||||
from langchain_core.tools import BaseTool
|
||||
from pydantic import PrivateAttr
|
||||
|
||||
from app.agent import StreamingCallbackHandler, ConversationMemoryManager
|
||||
from app.agent import StreamingHandler
|
||||
from app.chain import ChainBase
|
||||
from app.log import logger
|
||||
from app.schemas import Notification
|
||||
@@ -17,15 +16,16 @@ class ToolChain(ChainBase):
|
||||
|
||||
|
||||
class MoviePilotTool(BaseTool, metaclass=ABCMeta):
|
||||
"""MoviePilot专用工具基类"""
|
||||
"""
|
||||
MoviePilot专用工具基类(LangChain v1 / langchain_core)
|
||||
"""
|
||||
|
||||
_session_id: str = PrivateAttr()
|
||||
_user_id: str = PrivateAttr()
|
||||
_channel: str = PrivateAttr(default=None)
|
||||
_source: str = PrivateAttr(default=None)
|
||||
_username: str = PrivateAttr(default=None)
|
||||
_callback_handler: StreamingCallbackHandler = PrivateAttr(default=None)
|
||||
_memory_manager: ConversationMemoryManager = PrivateAttr(default=None)
|
||||
_channel: Optional[str] = PrivateAttr(default=None)
|
||||
_source: Optional[str] = PrivateAttr(default=None)
|
||||
_username: Optional[str] = PrivateAttr(default=None)
|
||||
_stream_handler: Optional[StreamingHandler] = PrivateAttr(default=None)
|
||||
|
||||
def __init__(self, session_id: str, user_id: str, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
@@ -33,69 +33,70 @@ class MoviePilotTool(BaseTool, metaclass=ABCMeta):
|
||||
self._user_id = user_id
|
||||
|
||||
def _run(self, *args: Any, **kwargs: Any) -> Any:
|
||||
pass
|
||||
raise NotImplementedError("MoviePilotTool 只支持异步调用,请使用 _arun")
|
||||
|
||||
async def _arun(self, **kwargs) -> str:
|
||||
"""异步运行工具"""
|
||||
# 发送和记忆工具调用前的信息
|
||||
agent_message = await self._callback_handler.get_message()
|
||||
if agent_message:
|
||||
# 发送消息
|
||||
await self.send_tool_message(agent_message, title="MoviePilot助手")
|
||||
|
||||
# 记忆工具调用
|
||||
await self._memory_manager.add_memory(
|
||||
session_id=self._session_id,
|
||||
user_id=self._user_id,
|
||||
role="tool_call",
|
||||
content=agent_message,
|
||||
metadata={
|
||||
"call_id": self.__class__.__name__,
|
||||
"tool_name": self.__class__.__name__,
|
||||
"parameters": kwargs
|
||||
}
|
||||
async def _arun(self, *args: Any, **kwargs: Any) -> str:
|
||||
"""
|
||||
异步运行工具,负责:
|
||||
1. 在工具调用前将流式消息推送给用户
|
||||
2. 持久化工具调用记录到会话记忆
|
||||
3. 调用具体工具逻辑(子类实现的 execute 方法)
|
||||
4. 持久化工具结果到会话记忆
|
||||
"""
|
||||
# 获取工具调用前 Agent 已积累的流式文本
|
||||
agent_message = (
|
||||
self._stream_handler.take() if self._stream_handler else ""
|
||||
)
|
||||
|
||||
# 发送执行工具说明,优先使用工具自定义的提示消息,如果没有则使用 explanation
|
||||
# 获取工具执行提示消息
|
||||
tool_message = self.get_tool_message(**kwargs)
|
||||
if not tool_message:
|
||||
explanation = kwargs.get("explanation")
|
||||
if explanation:
|
||||
tool_message = explanation
|
||||
|
||||
# 合并 Agent 消息和工具执行消息后一起发送
|
||||
messages = []
|
||||
if agent_message:
|
||||
messages.append(agent_message)
|
||||
if tool_message:
|
||||
formatted_message = f"⚙️ => {tool_message}"
|
||||
await self.send_tool_message(formatted_message)
|
||||
messages.append(f"⚙️ => {tool_message}")
|
||||
|
||||
logger.debug(f'Executing tool {self.name} with args: {kwargs}')
|
||||
result = await self.run(**kwargs)
|
||||
logger.debug(f'Tool {self.name} executed with result: {result}')
|
||||
if messages:
|
||||
merged_message = "\n\n".join(messages)
|
||||
await self.send_tool_message(merged_message, title="MoviePilot助手")
|
||||
|
||||
# 记忆工具调用结果
|
||||
logger.debug(f"Executing tool {self.name} with args: {kwargs}")
|
||||
|
||||
# 执行具体工具逻辑
|
||||
try:
|
||||
result = await self.run(**kwargs)
|
||||
logger.debug(f"Tool {self.name} executed with result: {result}")
|
||||
except Exception as e:
|
||||
error_message = f"工具执行异常 ({type(e).__name__}): {str(e)}"
|
||||
logger.error(f"Tool {self.name} execution failed: {e}", exc_info=True)
|
||||
result = error_message
|
||||
|
||||
# 格式化结果
|
||||
if isinstance(result, str):
|
||||
formated_result = result
|
||||
elif isinstance(result, int, float):
|
||||
formated_result = str(result)
|
||||
formatted_result = result
|
||||
elif isinstance(result, (int, float)):
|
||||
formatted_result = str(result)
|
||||
else:
|
||||
formated_result = json.dumps(result, ensure_ascii=False, indent=2)
|
||||
await self._memory_manager.add_memory(
|
||||
session_id=self._session_id,
|
||||
user_id=self._user_id,
|
||||
role="tool_result",
|
||||
content=formated_result
|
||||
)
|
||||
formatted_result = json.dumps(result, ensure_ascii=False, indent=2)
|
||||
|
||||
return result
|
||||
return formatted_result
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""
|
||||
获取工具执行时的友好提示消息
|
||||
|
||||
获取工具执行时的友好提示消息。
|
||||
|
||||
子类可以重写此方法,根据实际参数生成个性化的提示消息。
|
||||
如果返回 None 或空字符串,将回退使用 explanation 参数。
|
||||
|
||||
|
||||
Args:
|
||||
**kwargs: 工具的所有参数(包括 explanation)
|
||||
|
||||
|
||||
Returns:
|
||||
str: 友好的提示消息,如果返回 None 或空字符串则使用 explanation
|
||||
"""
|
||||
@@ -103,24 +104,27 @@ class MoviePilotTool(BaseTool, metaclass=ABCMeta):
|
||||
|
||||
@abstractmethod
|
||||
async def run(self, **kwargs) -> str:
|
||||
"""子类实现具体的工具执行逻辑"""
|
||||
raise NotImplementedError
|
||||
|
||||
def set_message_attr(self, channel: str, source: str, username: str):
|
||||
"""设置消息属性"""
|
||||
"""
|
||||
设置消息属性
|
||||
"""
|
||||
self._channel = channel
|
||||
self._source = source
|
||||
self._username = username
|
||||
|
||||
def set_callback_handler(self, callback_handler: StreamingCallbackHandler):
|
||||
"""设置回调处理器"""
|
||||
self._callback_handler = callback_handler
|
||||
|
||||
def set_memory_manager(self, memory_manager: ConversationMemoryManager):
|
||||
"""设置记忆客理器"""
|
||||
self._memory_manager = memory_manager
|
||||
def set_stream_handler(self, stream_handler: StreamingHandler):
|
||||
"""
|
||||
设置回调处理器
|
||||
"""
|
||||
self._stream_handler = stream_handler
|
||||
|
||||
async def send_tool_message(self, message: str, title: str = ""):
|
||||
"""发送工具消息"""
|
||||
"""
|
||||
发送工具消息
|
||||
"""
|
||||
await ToolChain().async_post_message(
|
||||
Notification(
|
||||
channel=self._channel,
|
||||
@@ -128,6 +132,6 @@ class MoviePilotTool(BaseTool, metaclass=ABCMeta):
|
||||
userid=self._user_id,
|
||||
username=self._username,
|
||||
title=title,
|
||||
text=message
|
||||
text=message,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
"""MoviePilot工具工厂"""
|
||||
|
||||
from typing import List, Callable
|
||||
|
||||
from app.agent.tools.impl.add_download import AddDownloadTool
|
||||
@@ -27,7 +25,9 @@ from app.agent.tools.impl.search_person_credits import SearchPersonCreditsTool
|
||||
from app.agent.tools.impl.recognize_media import RecognizeMediaTool
|
||||
from app.agent.tools.impl.scrape_metadata import ScrapeMetadataTool
|
||||
from app.agent.tools.impl.query_episode_schedule import QueryEpisodeScheduleTool
|
||||
from app.agent.tools.impl.query_media_detail import QueryMediaDetailTool
|
||||
from app.agent.tools.impl.search_torrents import SearchTorrentsTool
|
||||
from app.agent.tools.impl.get_search_results import GetSearchResultsTool
|
||||
from app.agent.tools.impl.search_web import SearchWebTool
|
||||
from app.agent.tools.impl.send_message import SendMessageTool
|
||||
from app.agent.tools.impl.query_schedulers import QuerySchedulersTool
|
||||
@@ -40,19 +40,27 @@ from app.agent.tools.impl.query_directory_settings import QueryDirectorySettings
|
||||
from app.agent.tools.impl.list_directory import ListDirectoryTool
|
||||
from app.agent.tools.impl.query_transfer_history import QueryTransferHistoryTool
|
||||
from app.agent.tools.impl.transfer_file import TransferFileTool
|
||||
from app.agent.tools.impl.execute_command import ExecuteCommandTool
|
||||
from app.agent.tools.impl.edit_file import EditFileTool
|
||||
from app.agent.tools.impl.write_file import WriteFileTool
|
||||
from app.agent.tools.impl.read_file import ReadFileTool
|
||||
from app.core.plugin import PluginManager
|
||||
from app.log import logger
|
||||
from .base import MoviePilotTool
|
||||
|
||||
|
||||
class MoviePilotToolFactory:
|
||||
"""MoviePilot工具工厂"""
|
||||
"""
|
||||
MoviePilot工具工厂
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def create_tools(session_id: str, user_id: str,
|
||||
channel: str = None, source: str = None, username: str = None,
|
||||
callback_handler: Callable = None, memory_mananger: Callable = None) -> List[MoviePilotTool]:
|
||||
"""创建MoviePilot工具列表"""
|
||||
stream_handler: Callable = None) -> List[MoviePilotTool]:
|
||||
"""
|
||||
创建MoviePilot工具列表
|
||||
"""
|
||||
tools = []
|
||||
tool_definitions = [
|
||||
SearchMediaTool,
|
||||
@@ -61,10 +69,12 @@ class MoviePilotToolFactory:
|
||||
RecognizeMediaTool,
|
||||
ScrapeMetadataTool,
|
||||
QueryEpisodeScheduleTool,
|
||||
QueryMediaDetailTool,
|
||||
AddSubscribeTool,
|
||||
UpdateSubscribeTool,
|
||||
SearchSubscribeTool,
|
||||
SearchTorrentsTool,
|
||||
GetSearchResultsTool,
|
||||
SearchWebTool,
|
||||
AddDownloadTool,
|
||||
QuerySubscribesTool,
|
||||
@@ -92,7 +102,11 @@ class MoviePilotToolFactory:
|
||||
QuerySchedulersTool,
|
||||
RunSchedulerTool,
|
||||
QueryWorkflowsTool,
|
||||
RunWorkflowTool
|
||||
RunWorkflowTool,
|
||||
ExecuteCommandTool,
|
||||
EditFileTool,
|
||||
WriteFileTool,
|
||||
ReadFileTool
|
||||
]
|
||||
# 创建内置工具
|
||||
for ToolClass in tool_definitions:
|
||||
@@ -101,8 +115,7 @@ class MoviePilotToolFactory:
|
||||
user_id=user_id
|
||||
)
|
||||
tool.set_message_attr(channel=channel, source=source, username=username)
|
||||
tool.set_callback_handler(callback_handler=callback_handler)
|
||||
tool.set_memory_manager(memory_manager=memory_mananger)
|
||||
tool.set_stream_handler(stream_handler=stream_handler)
|
||||
tools.append(tool)
|
||||
|
||||
# 加载插件提供的工具
|
||||
@@ -124,8 +137,7 @@ class MoviePilotToolFactory:
|
||||
user_id=user_id
|
||||
)
|
||||
tool.set_message_attr(channel=channel, source=source, username=username)
|
||||
tool.set_callback_handler(callback_handler=callback_handler)
|
||||
tool.set_memory_manager(memory_manager=memory_mananger)
|
||||
tool.set_stream_handler(stream_handler=stream_handler)
|
||||
tools.append(tool)
|
||||
plugin_tools_count += 1
|
||||
logger.debug(f"成功加载插件 {plugin_name}({plugin_id}) 的工具: {ToolClass.__name__}")
|
||||
|
||||
176
app/agent/tools/impl/_torrent_search_utils.py
Normal file
176
app/agent/tools/impl/_torrent_search_utils.py
Normal file
@@ -0,0 +1,176 @@
|
||||
"""种子搜索工具辅助函数"""
|
||||
|
||||
import re
|
||||
from typing import List, Optional
|
||||
|
||||
from app.core.context import Context
|
||||
from app.utils.crypto import HashUtils
|
||||
from app.utils.string import StringUtils
|
||||
|
||||
SEARCH_RESULT_CACHE_FILE = "__search_result__"
|
||||
TORRENT_RESULT_LIMIT = 50
|
||||
|
||||
|
||||
def build_torrent_ref(context: Optional[Context]) -> str:
|
||||
"""生成用于下载校验的短引用"""
|
||||
if not context or not context.torrent_info:
|
||||
return ""
|
||||
return HashUtils.sha1(context.torrent_info.enclosure or "")[:7]
|
||||
|
||||
|
||||
def sort_season_options(options: List[str]) -> List[str]:
|
||||
"""按前端逻辑排序季集选项"""
|
||||
if len(options) <= 1:
|
||||
return options
|
||||
|
||||
parsed_options = []
|
||||
for index, option in enumerate(options):
|
||||
match = re.match(r"^S(\d+)(?:-S(\d+))?\s*(?:E(\d+)(?:-E(\d+))?)?$", option or "")
|
||||
if not match:
|
||||
parsed_options.append({
|
||||
"original": option,
|
||||
"season_num": 0,
|
||||
"episode_num": 0,
|
||||
"max_episode_num": 0,
|
||||
"is_whole_season": False,
|
||||
"index": index,
|
||||
})
|
||||
continue
|
||||
|
||||
episode_num = int(match.group(3)) if match.group(3) else 0
|
||||
max_episode_num = int(match.group(4)) if match.group(4) else episode_num
|
||||
parsed_options.append({
|
||||
"original": option,
|
||||
"season_num": int(match.group(1)),
|
||||
"episode_num": episode_num,
|
||||
"max_episode_num": max_episode_num,
|
||||
"is_whole_season": not match.group(3),
|
||||
"index": index,
|
||||
})
|
||||
|
||||
whole_seasons = [item for item in parsed_options if item["is_whole_season"]]
|
||||
episodes = [item for item in parsed_options if not item["is_whole_season"]]
|
||||
|
||||
whole_seasons.sort(key=lambda item: (-item["season_num"], item["index"]))
|
||||
episodes.sort(
|
||||
key=lambda item: (
|
||||
-item["season_num"],
|
||||
-(item["max_episode_num"] or item["episode_num"]),
|
||||
-item["episode_num"],
|
||||
item["index"],
|
||||
)
|
||||
)
|
||||
return [item["original"] for item in whole_seasons + episodes]
|
||||
|
||||
|
||||
def append_option(options: List[str], value: Optional[str]) -> None:
|
||||
"""按前端逻辑收集去重后的筛选项"""
|
||||
if value and value not in options:
|
||||
options.append(value)
|
||||
|
||||
|
||||
def build_filter_options(items: List[Context]) -> dict:
|
||||
"""从搜索结果中构建筛选项汇总"""
|
||||
filter_options = {
|
||||
"site": [],
|
||||
"season": [],
|
||||
"freeState": [],
|
||||
"edition": [],
|
||||
"resolution": [],
|
||||
"videoCode": [],
|
||||
"releaseGroup": [],
|
||||
}
|
||||
|
||||
for item in items:
|
||||
torrent_info = item.torrent_info
|
||||
meta_info = item.meta_info
|
||||
append_option(filter_options["site"], getattr(torrent_info, "site_name", None))
|
||||
append_option(filter_options["season"], getattr(meta_info, "season_episode", None))
|
||||
append_option(filter_options["freeState"], getattr(torrent_info, "volume_factor", None))
|
||||
append_option(filter_options["edition"], getattr(meta_info, "edition", None))
|
||||
append_option(filter_options["resolution"], getattr(meta_info, "resource_pix", None))
|
||||
append_option(filter_options["videoCode"], getattr(meta_info, "video_encode", None))
|
||||
append_option(filter_options["releaseGroup"], getattr(meta_info, "resource_team", None))
|
||||
|
||||
filter_options["season"] = sort_season_options(filter_options["season"])
|
||||
return filter_options
|
||||
|
||||
|
||||
def match_filter(filter_values: Optional[List[str]], value: Optional[str]) -> bool:
|
||||
"""匹配前端同款多选筛选规则"""
|
||||
return not filter_values or bool(value and value in filter_values)
|
||||
|
||||
|
||||
def filter_contexts(items: List[Context],
|
||||
site: Optional[List[str]] = None,
|
||||
season: Optional[List[str]] = None,
|
||||
free_state: Optional[List[str]] = None,
|
||||
video_code: Optional[List[str]] = None,
|
||||
edition: Optional[List[str]] = None,
|
||||
resolution: Optional[List[str]] = None,
|
||||
release_group: Optional[List[str]] = None) -> List[Context]:
|
||||
"""按前端同款维度筛选结果"""
|
||||
filtered_items = []
|
||||
for item in items:
|
||||
torrent_info = item.torrent_info
|
||||
meta_info = item.meta_info
|
||||
if (
|
||||
match_filter(site, getattr(torrent_info, "site_name", None))
|
||||
and match_filter(free_state, getattr(torrent_info, "volume_factor", None))
|
||||
and match_filter(season, getattr(meta_info, "season_episode", None))
|
||||
and match_filter(release_group, getattr(meta_info, "resource_team", None))
|
||||
and match_filter(video_code, getattr(meta_info, "video_encode", None))
|
||||
and match_filter(resolution, getattr(meta_info, "resource_pix", None))
|
||||
and match_filter(edition, getattr(meta_info, "edition", None))
|
||||
):
|
||||
filtered_items.append(item)
|
||||
return filtered_items
|
||||
|
||||
|
||||
def simplify_search_result(context: Context, index: int) -> dict:
|
||||
"""精简单条搜索结果"""
|
||||
simplified = {}
|
||||
torrent_info = context.torrent_info
|
||||
meta_info = context.meta_info
|
||||
media_info = context.media_info
|
||||
|
||||
if torrent_info:
|
||||
simplified["torrent_info"] = {
|
||||
"title": torrent_info.title,
|
||||
"size": StringUtils.format_size(torrent_info.size),
|
||||
"seeders": torrent_info.seeders,
|
||||
"peers": torrent_info.peers,
|
||||
"site_name": torrent_info.site_name,
|
||||
"torrent_url": f"{build_torrent_ref(context)}:{index}",
|
||||
"page_url": torrent_info.page_url,
|
||||
"volume_factor": torrent_info.volume_factor,
|
||||
"freedate_diff": torrent_info.freedate_diff,
|
||||
"pubdate": torrent_info.pubdate,
|
||||
}
|
||||
|
||||
if media_info:
|
||||
simplified["media_info"] = {
|
||||
"title": media_info.title,
|
||||
"en_title": media_info.en_title,
|
||||
"year": media_info.year,
|
||||
"type": media_info.type.value if media_info.type else None,
|
||||
"season": media_info.season,
|
||||
"tmdb_id": media_info.tmdb_id,
|
||||
}
|
||||
|
||||
if meta_info:
|
||||
simplified["meta_info"] = {
|
||||
"name": meta_info.name,
|
||||
"cn_name": meta_info.cn_name,
|
||||
"en_name": meta_info.en_name,
|
||||
"year": meta_info.year,
|
||||
"type": meta_info.type.value if meta_info.type else None,
|
||||
"begin_season": meta_info.begin_season,
|
||||
"season_episode": meta_info.season_episode,
|
||||
"resource_team": meta_info.resource_team,
|
||||
"video_encode": meta_info.video_encode,
|
||||
"edition": meta_info.edition,
|
||||
"resource_pix": meta_info.resource_pix,
|
||||
}
|
||||
|
||||
return simplified
|
||||
@@ -1,27 +1,31 @@
|
||||
"""添加下载工具"""
|
||||
|
||||
from typing import Optional, Type
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Type
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool, ToolChain
|
||||
from app.chain.search import SearchChain
|
||||
from app.chain.download import DownloadChain
|
||||
from app.core.config import settings
|
||||
from app.core.context import Context
|
||||
from app.core.metainfo import MetaInfo
|
||||
from app.db.site_oper import SiteOper
|
||||
from app.helper.directory import DirectoryHelper
|
||||
from app.log import logger
|
||||
from app.schemas import TorrentInfo
|
||||
from app.schemas import TorrentInfo, FileURI
|
||||
from app.utils.crypto import HashUtils
|
||||
|
||||
|
||||
class AddDownloadInput(BaseModel):
|
||||
"""添加下载工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
site_name: str = Field(..., description="Name of the torrent site/source (e.g., 'The Pirate Bay')")
|
||||
torrent_title: str = Field(...,
|
||||
description="The display name/title of the torrent (e.g., 'The.Matrix.1999.1080p.BluRay.x264')")
|
||||
torrent_url: str = Field(..., description="Direct URL to the torrent file (.torrent) or magnet link")
|
||||
torrent_description: Optional[str] = Field(None,
|
||||
description="Brief description of the torrent content (optional)")
|
||||
torrent_url: List[str] = Field(
|
||||
...,
|
||||
description="One or more torrent_url values. Supports refs from get_search_results (`hash:id`) and magnet links."
|
||||
)
|
||||
downloader: Optional[str] = Field(None,
|
||||
description="Name of the downloader to use (optional, uses default if not specified)")
|
||||
save_path: Optional[str] = Field(None,
|
||||
@@ -32,75 +36,242 @@ class AddDownloadInput(BaseModel):
|
||||
|
||||
class AddDownloadTool(MoviePilotTool):
|
||||
name: str = "add_download"
|
||||
description: str = "Add torrent download task to the configured downloader (qBittorrent, Transmission, etc.). Downloads the torrent file and starts the download process with specified settings."
|
||||
description: str = "Add torrent download tasks using refs from get_search_results or magnet links."
|
||||
args_schema: Type[BaseModel] = AddDownloadInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据下载参数生成友好的提示消息"""
|
||||
torrent_title = kwargs.get("torrent_title", "")
|
||||
site_name = kwargs.get("site_name", "")
|
||||
torrent_urls = self._normalize_torrent_urls(kwargs.get("torrent_url"))
|
||||
downloader = kwargs.get("downloader")
|
||||
|
||||
message = f"正在添加下载任务: {torrent_title}"
|
||||
if site_name:
|
||||
message += f" (来源: {site_name})"
|
||||
|
||||
if torrent_urls:
|
||||
if len(torrent_urls) == 1:
|
||||
if self._is_torrent_ref(torrent_urls[0]):
|
||||
message = f"正在添加下载任务: 资源 {torrent_urls[0]}"
|
||||
else:
|
||||
message = "正在添加下载任务: 磁力链接"
|
||||
else:
|
||||
message = f"正在批量添加下载任务: 共 {len(torrent_urls)} 个资源"
|
||||
else:
|
||||
message = "正在添加下载任务"
|
||||
if downloader:
|
||||
message += f" [下载器: {downloader}]"
|
||||
|
||||
|
||||
return message
|
||||
|
||||
async def run(self, site_name: str, torrent_title: str, torrent_url: str, torrent_description: Optional[str] = None,
|
||||
@staticmethod
|
||||
def _build_torrent_ref(context: Context) -> str:
|
||||
"""生成用于校验缓存项的短引用"""
|
||||
if not context or not context.torrent_info:
|
||||
return ""
|
||||
return HashUtils.sha1(context.torrent_info.enclosure or "")[:7]
|
||||
|
||||
@staticmethod
|
||||
def _is_torrent_ref(torrent_ref: Optional[str]) -> bool:
|
||||
"""判断是否为内部搜索结果引用"""
|
||||
if not torrent_ref:
|
||||
return False
|
||||
return bool(re.fullmatch(r"[0-9a-f]{7}:\d+", str(torrent_ref).strip()))
|
||||
|
||||
@staticmethod
|
||||
def _is_magnet_link_input(torrent_url: Optional[str]) -> bool:
|
||||
"""判断输入是否为允许直接添加的磁力链接"""
|
||||
if not torrent_url:
|
||||
return False
|
||||
value = str(torrent_url).strip()
|
||||
return value.startswith("magnet:")
|
||||
|
||||
@classmethod
|
||||
def _resolve_cached_context(cls, torrent_ref: str) -> Optional[Context]:
|
||||
"""从最近一次搜索缓存中解析种子上下文,仅支持 hash:id 格式"""
|
||||
ref = str(torrent_ref).strip()
|
||||
if ":" not in ref:
|
||||
return None
|
||||
try:
|
||||
ref_hash, ref_index = ref.split(":", 1)
|
||||
index = int(ref_index)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
if index < 1:
|
||||
return None
|
||||
|
||||
results = SearchChain().last_search_results() or []
|
||||
if index > len(results):
|
||||
return None
|
||||
context = results[index - 1]
|
||||
if not ref_hash or cls._build_torrent_ref(context) != ref_hash:
|
||||
return None
|
||||
return context
|
||||
|
||||
@staticmethod
|
||||
def _merge_labels_with_system_tag(labels: Optional[str]) -> Optional[str]:
|
||||
"""合并用户标签与系统默认标签,确保任务可被系统管理"""
|
||||
system_tag = (settings.TORRENT_TAG or "").strip()
|
||||
user_labels = [item.strip() for item in (labels or "").split(",") if item.strip()]
|
||||
|
||||
if system_tag and system_tag not in user_labels:
|
||||
user_labels.append(system_tag)
|
||||
|
||||
return ",".join(user_labels) if user_labels else None
|
||||
|
||||
@staticmethod
|
||||
def _format_failed_result(failed_messages: List[str]) -> str:
|
||||
"""统一格式化失败结果"""
|
||||
return ", ".join([message for message in failed_messages if message])
|
||||
|
||||
@staticmethod
|
||||
def _build_failure_message(torrent_ref: str, error_msg: Optional[str] = None) -> str:
|
||||
"""构造失败提示"""
|
||||
normalized_error = (error_msg or "").strip()
|
||||
prefix = "添加种子任务失败:"
|
||||
if normalized_error.startswith(prefix):
|
||||
normalized_error = normalized_error[len(prefix):].lstrip()
|
||||
if AddDownloadTool._is_magnet_link_input(normalized_error):
|
||||
normalized_error = ""
|
||||
if normalized_error:
|
||||
return f"{torrent_ref} {normalized_error}"
|
||||
if AddDownloadTool._is_torrent_ref(torrent_ref):
|
||||
return torrent_ref
|
||||
return ""
|
||||
|
||||
@classmethod
|
||||
def _normalize_torrent_urls(cls, torrent_url: Optional[List[str] | str]) -> List[str]:
|
||||
"""统一规范 torrent_url 输入,保留所有非空值"""
|
||||
if torrent_url is None:
|
||||
return []
|
||||
|
||||
if isinstance(torrent_url, str):
|
||||
candidates = torrent_url.split(",")
|
||||
else:
|
||||
candidates = torrent_url
|
||||
|
||||
return [str(item).strip() for item in candidates if item and str(item).strip()]
|
||||
|
||||
@staticmethod
|
||||
def _resolve_direct_download_dir(save_path: Optional[str]) -> Optional[Path]:
|
||||
"""解析直接下载使用的目录,优先使用 save_path,其次使用默认下载目录"""
|
||||
if save_path:
|
||||
return Path(save_path)
|
||||
|
||||
download_dirs = DirectoryHelper().get_download_dirs()
|
||||
if not download_dirs:
|
||||
return None
|
||||
|
||||
dir_conf = download_dirs[0]
|
||||
if not dir_conf.download_path:
|
||||
return None
|
||||
|
||||
return Path(FileURI(storage=dir_conf.storage or "local", path=dir_conf.download_path).uri)
|
||||
|
||||
async def run(self, torrent_url: Optional[List[str]] = None,
|
||||
downloader: Optional[str] = None, save_path: Optional[str] = None,
|
||||
labels: Optional[str] = None, **kwargs) -> str:
|
||||
logger.info(
|
||||
f"执行工具: {self.name}, 参数: site_name={site_name}, torrent_title={torrent_title}, torrent_url={torrent_url}, downloader={downloader}, save_path={save_path}, labels={labels}")
|
||||
f"执行工具: {self.name}, 参数: torrent_url={torrent_url}, downloader={downloader}, save_path={save_path}, labels={labels}")
|
||||
|
||||
try:
|
||||
if not torrent_title or not torrent_url:
|
||||
return "错误:必须提供种子标题和下载链接"
|
||||
torrent_inputs = self._normalize_torrent_urls(torrent_url)
|
||||
if not torrent_inputs:
|
||||
return "错误:torrent_url 不能为空。"
|
||||
|
||||
# 使用DownloadChain添加下载
|
||||
download_chain = DownloadChain()
|
||||
merged_labels = self._merge_labels_with_system_tag(labels)
|
||||
success_count = 0
|
||||
failed_messages = []
|
||||
|
||||
# 根据站点名称查询站点cookie
|
||||
if not site_name:
|
||||
return "错误:必须提供站点名称,请从搜索资源结果信息中获取"
|
||||
siteinfo = await SiteOper().async_get_by_name(site_name)
|
||||
if not siteinfo:
|
||||
return f"错误:未找到站点信息:{site_name}"
|
||||
for torrent_input in torrent_inputs:
|
||||
if self._is_torrent_ref(torrent_input):
|
||||
cached_context = self._resolve_cached_context(torrent_input)
|
||||
if not cached_context or not cached_context.torrent_info:
|
||||
failed_messages.append(f"{torrent_input} 引用无效,请重新使用 get_search_results 查看搜索结果")
|
||||
continue
|
||||
|
||||
# 创建下载上下文
|
||||
torrent_info = TorrentInfo(
|
||||
title=torrent_title,
|
||||
description=torrent_description,
|
||||
enclosure=torrent_url,
|
||||
site_name=site_name,
|
||||
site_ua=siteinfo.ua,
|
||||
site_cookie=siteinfo.cookie,
|
||||
site_proxy=siteinfo.proxy,
|
||||
site_order=siteinfo.pri,
|
||||
site_downloader=siteinfo.downloader
|
||||
)
|
||||
meta_info = MetaInfo(title=torrent_title, subtitle=torrent_description)
|
||||
media_info = await ToolChain().async_recognize_media(meta=meta_info)
|
||||
if not media_info:
|
||||
return "错误:无法识别媒体信息,无法添加下载任务"
|
||||
context = Context(
|
||||
torrent_info=torrent_info,
|
||||
meta_info=meta_info,
|
||||
media_info=media_info
|
||||
)
|
||||
cached_torrent = cached_context.torrent_info
|
||||
site_name = cached_torrent.site_name
|
||||
torrent_title = cached_torrent.title or torrent_input
|
||||
torrent_description = cached_torrent.description
|
||||
enclosure = cached_torrent.enclosure
|
||||
|
||||
did = download_chain.download_single(
|
||||
context=context,
|
||||
downloader=downloader,
|
||||
save_path=save_path,
|
||||
label=labels
|
||||
)
|
||||
if did:
|
||||
return f"成功添加下载任务:{torrent_title}"
|
||||
else:
|
||||
return "添加下载任务失败"
|
||||
if not site_name:
|
||||
failed_messages.append(f"{torrent_input} 缺少站点名称")
|
||||
continue
|
||||
|
||||
siteinfo = await SiteOper().async_get_by_name(site_name)
|
||||
if not siteinfo:
|
||||
failed_messages.append(f"{torrent_input} 未找到站点信息 {site_name}")
|
||||
continue
|
||||
|
||||
torrent_info = TorrentInfo(
|
||||
title=torrent_title,
|
||||
description=torrent_description,
|
||||
enclosure=enclosure,
|
||||
site_name=site_name,
|
||||
site_ua=siteinfo.ua,
|
||||
site_cookie=siteinfo.cookie,
|
||||
site_proxy=siteinfo.proxy,
|
||||
site_order=siteinfo.pri,
|
||||
site_downloader=siteinfo.downloader
|
||||
)
|
||||
meta_info = MetaInfo(title=torrent_title, subtitle=torrent_description)
|
||||
media_info = cached_context.media_info if cached_context.media_info else None
|
||||
if not media_info:
|
||||
media_info = await ToolChain().async_recognize_media(meta=meta_info)
|
||||
if not media_info:
|
||||
failed_messages.append(f"{torrent_input} 无法识别媒体信息")
|
||||
continue
|
||||
|
||||
context = Context(
|
||||
torrent_info=torrent_info,
|
||||
meta_info=meta_info,
|
||||
media_info=media_info
|
||||
)
|
||||
else:
|
||||
if not self._is_magnet_link_input(torrent_input):
|
||||
failed_messages.append(
|
||||
f"{torrent_input} 不是有效的下载内容,非 hash:id 时仅支持 magnet: 开头"
|
||||
)
|
||||
continue
|
||||
download_dir = self._resolve_direct_download_dir(save_path)
|
||||
if not download_dir:
|
||||
failed_messages.append(f"{torrent_input} 缺少保存路径,且系统未配置可用下载目录")
|
||||
continue
|
||||
result = download_chain.download(
|
||||
content=torrent_input,
|
||||
download_dir=download_dir,
|
||||
cookie=None,
|
||||
label=merged_labels,
|
||||
downloader=downloader
|
||||
)
|
||||
if result:
|
||||
_, did, _, error_msg = result
|
||||
else:
|
||||
did, error_msg = None, "未找到下载器"
|
||||
if did:
|
||||
success_count += 1
|
||||
else:
|
||||
failed_messages.append(self._build_failure_message(torrent_input, error_msg))
|
||||
continue
|
||||
|
||||
did, error_msg = download_chain.download_single(
|
||||
context=context,
|
||||
downloader=downloader,
|
||||
save_path=save_path,
|
||||
label=merged_labels,
|
||||
return_detail=True
|
||||
)
|
||||
if did:
|
||||
success_count += 1
|
||||
else:
|
||||
failed_messages.append(self._build_failure_message(torrent_input, error_msg))
|
||||
|
||||
if success_count and not failed_messages:
|
||||
return "任务添加成功"
|
||||
|
||||
if success_count:
|
||||
return f"部分任务添加失败:{self._format_failed_result(failed_messages)}"
|
||||
|
||||
return f"任务添加失败:{self._format_failed_result(failed_messages)}"
|
||||
except Exception as e:
|
||||
logger.error(f"添加下载任务失败: {e}", exc_info=True)
|
||||
return f"添加下载任务时发生错误: {str(e)}"
|
||||
|
||||
@@ -16,11 +16,13 @@ class AddSubscribeInput(BaseModel):
|
||||
title: str = Field(..., description="The title of the media to subscribe to (e.g., 'The Matrix', 'Breaking Bad')")
|
||||
year: str = Field(..., description="Release year of the media (required for accurate identification)")
|
||||
media_type: str = Field(...,
|
||||
description="Type of media content: '电影' for films, '电视剧' for television series or anime series")
|
||||
description="Allowed values: movie, tv")
|
||||
season: Optional[int] = Field(None,
|
||||
description="Season number for TV shows (optional, if not specified will subscribe to all seasons)")
|
||||
tmdb_id: Optional[str] = Field(None,
|
||||
description="TMDB database ID for precise media identification (optional but recommended for accuracy)")
|
||||
tmdb_id: Optional[int] = Field(None,
|
||||
description="TMDB database ID for precise media identification (optional, can be obtained from search_media tool)")
|
||||
douban_id: Optional[str] = Field(None,
|
||||
description="Douban ID for precise media identification (optional, alternative to tmdb_id)")
|
||||
start_episode: Optional[int] = Field(None,
|
||||
description="Starting episode number for TV shows (optional, defaults to 1 if not specified)")
|
||||
total_episode: Optional[int] = Field(None,
|
||||
@@ -32,9 +34,9 @@ class AddSubscribeInput(BaseModel):
|
||||
effect: Optional[str] = Field(None,
|
||||
description="Effect filter as regular expression (optional, e.g., 'HDR|DV|SDR')")
|
||||
filter_groups: Optional[List[str]] = Field(None,
|
||||
description="List of filter rule group names to apply (optional, use query_rule_groups tool to get available rule groups)")
|
||||
description="List of filter rule group names to apply (optional, can be obtained from query_rule_groups tool)")
|
||||
sites: Optional[List[int]] = Field(None,
|
||||
description="List of site IDs to search from (optional, use query_sites tool to get available site IDs)")
|
||||
description="List of site IDs to search from (optional, can be obtained from query_sites tool)")
|
||||
|
||||
|
||||
class AddSubscribeTool(MoviePilotTool):
|
||||
@@ -60,26 +62,23 @@ class AddSubscribeTool(MoviePilotTool):
|
||||
return message
|
||||
|
||||
async def run(self, title: str, year: str, media_type: str,
|
||||
season: Optional[int] = None, tmdb_id: Optional[str] = None,
|
||||
season: Optional[int] = None, tmdb_id: Optional[int] = None,
|
||||
douban_id: Optional[str] = None,
|
||||
start_episode: Optional[int] = None, total_episode: Optional[int] = None,
|
||||
quality: Optional[str] = None, resolution: Optional[str] = None,
|
||||
effect: Optional[str] = None, filter_groups: Optional[List[str]] = None,
|
||||
sites: Optional[List[int]] = None, **kwargs) -> str:
|
||||
logger.info(
|
||||
f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, "
|
||||
f"season={season}, tmdb_id={tmdb_id}, start_episode={start_episode}, "
|
||||
f"season={season}, tmdb_id={tmdb_id}, douban_id={douban_id}, start_episode={start_episode}, "
|
||||
f"total_episode={total_episode}, quality={quality}, resolution={resolution}, "
|
||||
f"effect={effect}, filter_groups={filter_groups}, sites={sites}")
|
||||
|
||||
try:
|
||||
subscribe_chain = SubscribeChain()
|
||||
# 转换 tmdb_id 为整数
|
||||
tmdbid_int = None
|
||||
if tmdb_id:
|
||||
try:
|
||||
tmdbid_int = int(tmdb_id)
|
||||
except (ValueError, TypeError):
|
||||
logger.warning(f"无效的 tmdb_id: {tmdb_id},将忽略")
|
||||
media_type_enum = MediaType.from_agent(media_type)
|
||||
if not media_type_enum:
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv'"
|
||||
|
||||
# 构建额外的订阅参数
|
||||
subscribe_kwargs = {}
|
||||
@@ -99,15 +98,19 @@ class AddSubscribeTool(MoviePilotTool):
|
||||
subscribe_kwargs['sites'] = sites
|
||||
|
||||
sid, message = await subscribe_chain.async_add(
|
||||
mtype=MediaType(media_type),
|
||||
mtype=media_type_enum,
|
||||
title=title,
|
||||
year=year,
|
||||
tmdbid=tmdbid_int,
|
||||
tmdbid=tmdb_id,
|
||||
doubanid=douban_id,
|
||||
season=season,
|
||||
username=self._user_id,
|
||||
**subscribe_kwargs
|
||||
)
|
||||
if sid:
|
||||
if message and "已存在" in message:
|
||||
return f"订阅已存在:{title} ({year})。如需修改参数请先删除旧订阅。"
|
||||
|
||||
result_msg = f"成功添加订阅:{title} ({year})"
|
||||
if subscribe_kwargs:
|
||||
params = []
|
||||
|
||||
@@ -12,23 +12,23 @@ from app.log import logger
|
||||
class DeleteDownloadInput(BaseModel):
|
||||
"""删除下载任务工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
task_identifier: str = Field(..., description="Task identifier: can be task hash (unique identifier) or task title/name")
|
||||
hash: str = Field(..., description="Task hash (can be obtained from query_download_tasks tool)")
|
||||
downloader: Optional[str] = Field(None, description="Name of specific downloader (optional, if not provided will search all downloaders)")
|
||||
delete_files: Optional[bool] = Field(False, description="Whether to delete downloaded files along with the task (default: False, only removes the task from downloader)")
|
||||
|
||||
|
||||
class DeleteDownloadTool(MoviePilotTool):
|
||||
name: str = "delete_download"
|
||||
description: str = "Delete a download task from the downloader. Can delete by task hash (unique identifier) or task title/name. Optionally specify the downloader name and whether to delete downloaded files."
|
||||
description: str = "Delete a download task from the downloader by task hash only. Optionally specify the downloader name and whether to delete downloaded files."
|
||||
args_schema: Type[BaseModel] = DeleteDownloadInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据删除参数生成友好的提示消息"""
|
||||
task_identifier = kwargs.get("task_identifier", "")
|
||||
hash_value = kwargs.get("hash", "")
|
||||
downloader = kwargs.get("downloader")
|
||||
delete_files = kwargs.get("delete_files", False)
|
||||
|
||||
message = f"正在删除下载任务: {task_identifier}"
|
||||
message = f"正在删除下载任务: {hash_value}"
|
||||
if downloader:
|
||||
message += f" [下载器: {downloader}]"
|
||||
if delete_files:
|
||||
@@ -36,40 +36,26 @@ class DeleteDownloadTool(MoviePilotTool):
|
||||
|
||||
return message
|
||||
|
||||
async def run(self, task_identifier: str, downloader: Optional[str] = None,
|
||||
async def run(self, hash: str, downloader: Optional[str] = None,
|
||||
delete_files: Optional[bool] = False, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: task_identifier={task_identifier}, downloader={downloader}, delete_files={delete_files}")
|
||||
logger.info(f"执行工具: {self.name}, 参数: hash={hash}, downloader={downloader}, delete_files={delete_files}")
|
||||
|
||||
try:
|
||||
download_chain = DownloadChain()
|
||||
|
||||
# 如果task_identifier看起来像hash(通常是40个字符的十六进制字符串)
|
||||
task_hash = None
|
||||
if len(task_identifier) == 40 and all(c in '0123456789abcdefABCDEF' for c in task_identifier):
|
||||
# 直接使用hash
|
||||
task_hash = task_identifier
|
||||
else:
|
||||
# 通过标题查找任务
|
||||
downloads = download_chain.downloading(name=downloader)
|
||||
for dl in downloads:
|
||||
# 检查标题或名称是否匹配
|
||||
if (task_identifier.lower() in (dl.title or "").lower()) or \
|
||||
(task_identifier.lower() in (dl.name or "").lower()):
|
||||
task_hash = dl.hash
|
||||
break
|
||||
|
||||
if not task_hash:
|
||||
return f"未找到匹配的下载任务:{task_identifier},请使用 query_downloads 工具查询可用的下载任务"
|
||||
|
||||
# 仅支持通过hash删除任务
|
||||
if len(hash) != 40 or not all(c in '0123456789abcdefABCDEF' for c in hash):
|
||||
return "参数错误:hash 格式无效,请先使用 query_download_tasks 工具获取正确的 hash。"
|
||||
|
||||
# 删除下载任务
|
||||
# remove_torrents 支持 delete_file 参数,可以控制是否删除文件
|
||||
result = download_chain.remove_torrents(hashs=[task_hash], downloader=downloader, delete_file=delete_files)
|
||||
result = download_chain.remove_torrents(hashs=[hash], downloader=downloader, delete_file=delete_files)
|
||||
|
||||
if result:
|
||||
files_info = "(包含文件)" if delete_files else "(不包含文件)"
|
||||
return f"成功删除下载任务:{task_identifier} {files_info}"
|
||||
return f"成功删除下载任务:{hash} {files_info}"
|
||||
else:
|
||||
return f"删除下载任务失败:{task_identifier},请检查任务是否存在或下载器是否可用"
|
||||
return f"删除下载任务失败:{hash},请检查任务是否存在或下载器是否可用"
|
||||
except Exception as e:
|
||||
logger.error(f"删除下载任务失败: {e}", exc_info=True)
|
||||
return f"删除下载任务时发生错误: {str(e)}"
|
||||
|
||||
75
app/agent/tools/impl/edit_file.py
Normal file
75
app/agent/tools/impl/edit_file.py
Normal file
@@ -0,0 +1,75 @@
|
||||
"""文件编辑工具"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Optional, Type
|
||||
|
||||
from anyio import Path as AsyncPath
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.log import logger
|
||||
|
||||
|
||||
class EditFileInput(BaseModel):
|
||||
"""Input parameters for edit file tool"""
|
||||
file_path: str = Field(..., description="The absolute path of the file to edit")
|
||||
old_text: str = Field(..., description="The exact old text to be replaced")
|
||||
new_text: str = Field(..., description="The new text to replace with")
|
||||
|
||||
|
||||
class EditFileTool(MoviePilotTool):
|
||||
name: str = "edit_file"
|
||||
description: str = "Edit a file by replacing specific old text with new text. Useful for modifying configuration files, code, or scripts."
|
||||
args_schema: Type[BaseModel] = EditFileInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据参数生成友好的提示消息"""
|
||||
file_path = kwargs.get("file_path", "")
|
||||
file_name = Path(file_path).name if file_path else "未知文件"
|
||||
return f"正在编辑文件: {file_name}"
|
||||
|
||||
async def run(self, file_path: str, old_text: str, new_text: str, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: file_path={file_path}")
|
||||
|
||||
try:
|
||||
path = AsyncPath(file_path)
|
||||
# 校验逻辑:如果要替换特定文本,文件必须存在且包含该文本
|
||||
if not await path.exists():
|
||||
# 如果 old_text 为空,可能用户想直接创建文件,但通常 edit_file 需要匹配旧内容
|
||||
if old_text:
|
||||
return f"错误:文件 {file_path} 不存在,无法进行内容替换。"
|
||||
|
||||
if await path.exists() and not await path.is_file():
|
||||
return f"错误:{file_path} 不是一个文件"
|
||||
|
||||
if await path.exists():
|
||||
content = await path.read_text(encoding="utf-8")
|
||||
if old_text not in content:
|
||||
logger.warning(f"编辑文件 {file_path} 失败:未找到指定的旧文本块")
|
||||
return f"错误:在文件 {file_path} 中未找到指定的旧文本。请确保包含所有的空格、缩进 and 换行符。"
|
||||
occurrences = content.count(old_text)
|
||||
new_content = content.replace(old_text, new_text)
|
||||
else:
|
||||
# 文件不存在且 old_text 为空的情形(初始化新文件)
|
||||
new_content = new_text
|
||||
occurrences = 1
|
||||
|
||||
# 自动创建父目录
|
||||
await path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 写入文件
|
||||
await path.write_text(new_content, encoding="utf-8")
|
||||
|
||||
logger.info(f"成功编辑文件 {file_path},替换了 {occurrences} 处内容")
|
||||
return f"成功编辑文件 {file_path} (替换了 {occurrences} 处匹配内容)"
|
||||
|
||||
|
||||
except PermissionError:
|
||||
return f"错误:没有访问/修改 {file_path} 的权限"
|
||||
except UnicodeDecodeError:
|
||||
return f"错误:{file_path} 不是文本文件,无法编辑"
|
||||
except Exception as e:
|
||||
logger.error(f"编辑文件 {file_path} 时发生错误: {str(e)}", exc_info=True)
|
||||
return f"操作失败: {str(e)}"
|
||||
|
||||
|
||||
81
app/agent/tools/impl/execute_command.py
Normal file
81
app/agent/tools/impl/execute_command.py
Normal file
@@ -0,0 +1,81 @@
|
||||
"""执行Shell命令工具"""
|
||||
|
||||
import asyncio
|
||||
from typing import Optional, Type
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.log import logger
|
||||
|
||||
|
||||
class ExecuteCommandInput(BaseModel):
|
||||
"""执行Shell命令工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this command is being executed")
|
||||
command: str = Field(..., description="The shell command to execute")
|
||||
timeout: Optional[int] = Field(60, description="Max execution time in seconds (default: 60)")
|
||||
|
||||
|
||||
class ExecuteCommandTool(MoviePilotTool):
|
||||
name: str = "execute_command"
|
||||
description: str = "Safely execute shell commands on the server. Useful for system maintenance, checking status, or running custom scripts. Includes timeout and output limits."
|
||||
args_schema: Type[BaseModel] = ExecuteCommandInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据命令生成友好的提示消息"""
|
||||
command = kwargs.get("command", "")
|
||||
return f"正在执行系统命令: {command}"
|
||||
|
||||
async def run(self, command: str, timeout: Optional[int] = 60, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: command={command}, timeout={timeout}")
|
||||
|
||||
# 简单安全过滤
|
||||
forbidden_keywords = ["rm -rf /", ":(){ :|:& };:", "dd if=/dev/zero", "mkfs", "reboot", "shutdown"]
|
||||
for keyword in forbidden_keywords:
|
||||
if keyword in command:
|
||||
return f"错误:命令包含禁止使用的关键字 '{keyword}'"
|
||||
|
||||
try:
|
||||
# 执行命令
|
||||
process = await asyncio.create_subprocess_shell(
|
||||
command,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
try:
|
||||
# 等待完成,带超时
|
||||
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
|
||||
|
||||
# 处理输出
|
||||
stdout_str = stdout.decode('utf-8', errors='replace').strip()
|
||||
stderr_str = stderr.decode('utf-8', errors='replace').strip()
|
||||
exit_code = process.returncode
|
||||
|
||||
result = f"命令执行完成 (退出码: {exit_code})"
|
||||
if stdout_str:
|
||||
result += f"\n\n标准输出:\n{stdout_str}"
|
||||
if stderr_str:
|
||||
result += f"\n\n错误输出:\n{stderr_str}"
|
||||
|
||||
# 如果没有输出
|
||||
if not stdout_str and not stderr_str:
|
||||
result += "\n\n(无输出内容)"
|
||||
|
||||
# 限制输出长度,防止上下文过长
|
||||
if len(result) > 3000:
|
||||
result = result[:3000] + "\n\n...(输出内容过长,已截断)"
|
||||
|
||||
return result
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
# 超时处理
|
||||
try:
|
||||
process.kill()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
return f"命令执行超时 (限制: {timeout}秒)"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"执行命令失败: {e}", exc_info=True)
|
||||
return f"执行命令时发生错误: {str(e)}"
|
||||
@@ -8,6 +8,7 @@ from pydantic import BaseModel, Field
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.recommend import RecommendChain
|
||||
from app.log import logger
|
||||
from app.schemas.types import MediaType, media_type_to_agent
|
||||
|
||||
|
||||
class GetRecommendationsInput(BaseModel):
|
||||
@@ -30,7 +31,7 @@ class GetRecommendationsInput(BaseModel):
|
||||
"'douban_tv_animation' for Douban popular animation, "
|
||||
"'bangumi_calendar' for Bangumi anime calendar")
|
||||
media_type: Optional[str] = Field("all",
|
||||
description="Type of media content: '电影' for films, '电视剧' for television series or anime series, 'all' for all types")
|
||||
description="Allowed values: movie, tv, all")
|
||||
limit: Optional[int] = Field(20,
|
||||
description="Maximum number of recommendations to return (default: 20, maximum: 100)")
|
||||
|
||||
@@ -75,6 +76,12 @@ class GetRecommendationsTool(MoviePilotTool):
|
||||
media_type: Optional[str] = "all", limit: Optional[int] = 20, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: source={source}, media_type={media_type}, limit={limit}")
|
||||
try:
|
||||
if media_type != "all":
|
||||
media_type_enum = MediaType.from_agent(media_type)
|
||||
if not media_type_enum:
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv', 'all'"
|
||||
media_type = media_type_enum.to_agent() # 归一化为 "movie"/"tv"
|
||||
|
||||
recommend_chain = RecommendChain()
|
||||
results = []
|
||||
if source == "tmdb_trending":
|
||||
@@ -149,7 +156,7 @@ class GetRecommendationsTool(MoviePilotTool):
|
||||
"title": r.get("title"),
|
||||
"en_title": r.get("en_title"),
|
||||
"year": r.get("year"),
|
||||
"type": r.get("type"),
|
||||
"type": media_type_to_agent(r.get("type")),
|
||||
"season": r.get("season"),
|
||||
"tmdb_id": r.get("tmdb_id"),
|
||||
"imdb_id": r.get("imdb_id"),
|
||||
|
||||
108
app/agent/tools/impl/get_search_results.py
Normal file
108
app/agent/tools/impl/get_search_results.py
Normal file
@@ -0,0 +1,108 @@
|
||||
"""获取搜索结果工具"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import List, Optional, Type
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.search import SearchChain
|
||||
from app.log import logger
|
||||
from ._torrent_search_utils import (
|
||||
TORRENT_RESULT_LIMIT,
|
||||
build_filter_options,
|
||||
filter_contexts,
|
||||
simplify_search_result,
|
||||
)
|
||||
|
||||
|
||||
class GetSearchResultsInput(BaseModel):
|
||||
"""获取搜索结果工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
site: Optional[List[str]] = Field(None, description="Site name filters")
|
||||
season: Optional[List[str]] = Field(None, description="Season or episode filters")
|
||||
free_state: Optional[List[str]] = Field(None, description="Promotion state filters")
|
||||
video_code: Optional[List[str]] = Field(None, description="Video codec filters")
|
||||
edition: Optional[List[str]] = Field(None, description="Edition filters")
|
||||
resolution: Optional[List[str]] = Field(None, description="Resolution filters")
|
||||
release_group: Optional[List[str]] = Field(None, description="Release group filters")
|
||||
title_pattern: Optional[str] = Field(None, description="Regular expression pattern to filter torrent titles (e.g., '4K|2160p|UHD', '1080p.*BluRay')")
|
||||
show_filter_options: Optional[bool] = Field(False, description="Whether to return only optional filter options for re-checking available conditions")
|
||||
|
||||
class GetSearchResultsTool(MoviePilotTool):
|
||||
name: str = "get_search_results"
|
||||
description: str = "Get cached torrent search results from search_torrents with optional filters. Returns at most the first 50 matches."
|
||||
args_schema: Type[BaseModel] = GetSearchResultsInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
return "正在获取搜索结果"
|
||||
|
||||
async def run(self, site: Optional[List[str]] = None, season: Optional[List[str]] = None,
|
||||
free_state: Optional[List[str]] = None, video_code: Optional[List[str]] = None,
|
||||
edition: Optional[List[str]] = None, resolution: Optional[List[str]] = None,
|
||||
release_group: Optional[List[str]] = None, title_pattern: Optional[str] = None,
|
||||
show_filter_options: bool = False,
|
||||
**kwargs) -> str:
|
||||
logger.info(
|
||||
f"执行工具: {self.name}, 参数: site={site}, season={season}, free_state={free_state}, video_code={video_code}, edition={edition}, resolution={resolution}, release_group={release_group}, title_pattern={title_pattern}, show_filter_options={show_filter_options}")
|
||||
|
||||
try:
|
||||
items = await SearchChain().async_last_search_results() or []
|
||||
if not items:
|
||||
return "没有可用的搜索结果,请先使用 search_torrents 搜索"
|
||||
|
||||
if show_filter_options:
|
||||
payload = {
|
||||
"total_count": len(items),
|
||||
"filter_options": build_filter_options(items),
|
||||
}
|
||||
return json.dumps(payload, ensure_ascii=False, indent=2)
|
||||
|
||||
regex_pattern = None
|
||||
if title_pattern:
|
||||
try:
|
||||
regex_pattern = re.compile(title_pattern, re.IGNORECASE)
|
||||
except re.error as e:
|
||||
logger.warning(f"正则表达式编译失败: {title_pattern}, 错误: {e}")
|
||||
return f"正则表达式格式错误: {str(e)}"
|
||||
|
||||
filtered_items = filter_contexts(
|
||||
items=items,
|
||||
site=site,
|
||||
season=season,
|
||||
free_state=free_state,
|
||||
video_code=video_code,
|
||||
edition=edition,
|
||||
resolution=resolution,
|
||||
release_group=release_group,
|
||||
)
|
||||
if regex_pattern:
|
||||
filtered_items = [
|
||||
item for item in filtered_items
|
||||
if item.torrent_info and item.torrent_info.title
|
||||
and regex_pattern.search(item.torrent_info.title)
|
||||
]
|
||||
if not filtered_items:
|
||||
return "没有符合筛选条件的搜索结果,请调整筛选条件"
|
||||
|
||||
total_count = len(filtered_items)
|
||||
filtered_ids = {id(item) for item in filtered_items}
|
||||
matched_indices = [index for index, item in enumerate(items, start=1) if id(item) in filtered_ids]
|
||||
limited_items = filtered_items[:TORRENT_RESULT_LIMIT]
|
||||
limited_indices = matched_indices[:TORRENT_RESULT_LIMIT]
|
||||
results = [
|
||||
simplify_search_result(item, index)
|
||||
for item, index in zip(limited_items, limited_indices)
|
||||
]
|
||||
payload = {
|
||||
"total_count": total_count,
|
||||
"results": results,
|
||||
}
|
||||
if total_count > TORRENT_RESULT_LIMIT:
|
||||
payload["message"] = f"搜索结果共找到 {total_count} 条,仅显示前 {TORRENT_RESULT_LIMIT} 条结果。"
|
||||
return json.dumps(payload, ensure_ascii=False, indent=2)
|
||||
except Exception as e:
|
||||
error_message = f"获取搜索结果失败: {str(e)}"
|
||||
logger.error(f"获取搜索结果失败: {e}", exc_info=True)
|
||||
return error_message
|
||||
@@ -24,7 +24,7 @@ class ListDirectoryInput(BaseModel):
|
||||
|
||||
class ListDirectoryTool(MoviePilotTool):
|
||||
name: str = "list_directory"
|
||||
description: str = "List actual files and folders in a file system directory (NOT configuration). Shows files and subdirectories with their names, types, sizes, and modification times. Returns up to 20 items and the total count if there are more items. Use 'query_directories' to query directory configuration settings."
|
||||
description: str = "List actual files and folders in a file system directory (NOT configuration). Shows files and subdirectories with their names, types, sizes, and modification times. Returns up to 20 items and the total count if there are more items. Use 'query_directory_settings' to query directory configuration settings."
|
||||
args_schema: Type[BaseModel] = ListDirectoryInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
|
||||
@@ -10,7 +10,7 @@ from app.chain.download import DownloadChain
|
||||
from app.db.downloadhistory_oper import DownloadHistoryOper
|
||||
from app.log import logger
|
||||
from app.schemas import TransferTorrent, DownloadingTorrent
|
||||
from app.schemas.types import TorrentStatus
|
||||
from app.schemas.types import TorrentStatus, media_type_to_agent
|
||||
|
||||
|
||||
class QueryDownloadTasksInput(BaseModel):
|
||||
@@ -29,7 +29,8 @@ class QueryDownloadTasksTool(MoviePilotTool):
|
||||
description: str = "Query download status and list download tasks. Can query all active downloads, or search for specific tasks by hash or title. Shows download progress, completion status, and task details from configured downloaders."
|
||||
args_schema: Type[BaseModel] = QueryDownloadTasksInput
|
||||
|
||||
def _get_all_torrents(self, download_chain: DownloadChain, downloader: Optional[str] = None) -> List[Union[TransferTorrent, DownloadingTorrent]]:
|
||||
@staticmethod
|
||||
def _get_all_torrents(download_chain: DownloadChain, downloader: Optional[str] = None) -> List[Union[TransferTorrent, DownloadingTorrent]]:
|
||||
"""
|
||||
查询所有状态的任务(包括下载中和已完成的任务)
|
||||
"""
|
||||
@@ -50,6 +51,18 @@ class QueryDownloadTasksTool(MoviePilotTool):
|
||||
|
||||
return all_torrents
|
||||
|
||||
@staticmethod
|
||||
def _format_progress(progress: Optional[float]) -> Optional[str]:
|
||||
"""
|
||||
将下载进度格式化为保留一位小数的百分比字符串
|
||||
"""
|
||||
try:
|
||||
if progress is None:
|
||||
return None
|
||||
return f"{float(progress):.1f}%"
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据查询参数生成友好的提示消息"""
|
||||
downloader = kwargs.get("downloader")
|
||||
@@ -197,7 +210,7 @@ class QueryDownloadTasksTool(MoviePilotTool):
|
||||
"year": d.year,
|
||||
"season_episode": d.season_episode,
|
||||
"size": d.size,
|
||||
"progress": d.progress,
|
||||
"progress": self._format_progress(d.progress),
|
||||
"state": d.state,
|
||||
"upspeed": d.upspeed,
|
||||
"dlspeed": d.dlspeed,
|
||||
@@ -207,7 +220,7 @@ class QueryDownloadTasksTool(MoviePilotTool):
|
||||
if d.media:
|
||||
simplified["media"] = {
|
||||
"tmdbid": d.media.get("tmdbid"),
|
||||
"type": d.media.get("type"),
|
||||
"type": media_type_to_agent(d.media.get("type")),
|
||||
"title": d.media.get("title"),
|
||||
"season": d.media.get("season"),
|
||||
"episode": d.media.get("episode")
|
||||
|
||||
@@ -6,23 +6,21 @@ from typing import Optional, Type
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.media import MediaChain
|
||||
from app.chain.tmdb import TmdbChain
|
||||
from app.log import logger
|
||||
from app.schemas import MediaType
|
||||
|
||||
|
||||
class QueryEpisodeScheduleInput(BaseModel):
|
||||
"""查询剧集上映时间工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
tmdb_id: int = Field(..., description="TMDB ID of the TV series")
|
||||
tmdb_id: int = Field(..., description="TMDB ID of the TV series (can be obtained from search_media tool)")
|
||||
season: int = Field(..., description="Season number to query")
|
||||
episode_group: Optional[str] = Field(None, description="Episode group ID (optional)")
|
||||
|
||||
|
||||
class QueryEpisodeScheduleTool(MoviePilotTool):
|
||||
name: str = "query_episode_schedule"
|
||||
description: str = "Query TV series episode air dates and schedule. Returns detailed information for each episode including air date, episode number, title, overview, and other metadata. Filters out episodes without air dates."
|
||||
description: str = "Query TV series episode air dates and schedule. Returns non-duplicated schedule fields, including episode list, air-date statistics, and per-episode metadata. Filters out episodes without air dates."
|
||||
args_schema: Type[BaseModel] = QueryEpisodeScheduleInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
@@ -41,12 +39,6 @@ class QueryEpisodeScheduleTool(MoviePilotTool):
|
||||
logger.info(f"执行工具: {self.name}, 参数: tmdb_id={tmdb_id}, season={season}, episode_group={episode_group}")
|
||||
|
||||
try:
|
||||
# 获取媒体信息(用于获取标题和海报)
|
||||
media_chain = MediaChain()
|
||||
mediainfo = await media_chain.async_recognize_media(tmdbid=tmdb_id, mtype=MediaType.TV)
|
||||
if not mediainfo:
|
||||
return f"未找到 TMDB ID {tmdb_id} 的媒体信息"
|
||||
|
||||
# 获取集列表
|
||||
tmdb_chain = TmdbChain()
|
||||
episodes = await tmdb_chain.async_tmdb_episodes(
|
||||
@@ -92,12 +84,7 @@ class QueryEpisodeScheduleTool(MoviePilotTool):
|
||||
episode_list.sort(key=lambda x: (x["air_date"] or "", x["episode_number"] or 0))
|
||||
|
||||
result = {
|
||||
"success": True,
|
||||
"tmdb_id": tmdb_id,
|
||||
"season": season,
|
||||
"episode_group": episode_group,
|
||||
"series_title": mediainfo.title if mediainfo else None,
|
||||
"series_poster": mediainfo.poster_path if mediainfo else None,
|
||||
"total_episodes": len(episodes),
|
||||
"episodes_with_air_date": len(episode_list),
|
||||
"episodes": episode_list
|
||||
|
||||
@@ -1,96 +1,176 @@
|
||||
"""查询媒体库工具"""
|
||||
|
||||
import json
|
||||
from typing import Optional, Type
|
||||
from collections import OrderedDict
|
||||
from typing import Optional, Type, Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.mediaserver import MediaServerChain
|
||||
from app.core.context import MediaInfo
|
||||
from app.helper.mediaserver import MediaServerHelper
|
||||
from app.log import logger
|
||||
from app.schemas.types import MediaType
|
||||
from app.schemas.types import MediaType, media_type_to_agent
|
||||
|
||||
|
||||
def _sort_seasons(seasons: Optional[dict]) -> dict:
|
||||
"""按季号、集号升序整理季集信息,保证输出稳定。"""
|
||||
if not seasons:
|
||||
return {}
|
||||
|
||||
def _sort_key(value):
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return str(value)
|
||||
|
||||
return OrderedDict(
|
||||
(season, sorted(episodes, key=_sort_key))
|
||||
for season, episodes in sorted(seasons.items(), key=lambda item: _sort_key(item[0]))
|
||||
)
|
||||
|
||||
|
||||
def _filter_regular_seasons(seasons: Optional[dict]) -> OrderedDict:
|
||||
"""仅保留正片季,忽略 season 0 等特殊季。"""
|
||||
sorted_seasons = _sort_seasons(seasons)
|
||||
regular_seasons = OrderedDict()
|
||||
for season, episodes in sorted_seasons.items():
|
||||
try:
|
||||
season_number = int(season)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
if season_number > 0:
|
||||
regular_seasons[season_number] = episodes
|
||||
return regular_seasons
|
||||
|
||||
|
||||
def _build_tv_server_result(existing_seasons: OrderedDict, total_seasons: OrderedDict) -> dict[str, Any]:
|
||||
"""构建单个服务器的电视剧存在性结果。"""
|
||||
seasons_result = OrderedDict()
|
||||
missing_seasons = []
|
||||
all_seasons = sorted(set(total_seasons.keys()) | set(existing_seasons.keys()))
|
||||
|
||||
for season in all_seasons:
|
||||
existing_episodes = existing_seasons.get(season, [])
|
||||
total_episodes = total_seasons.get(season)
|
||||
if total_episodes is not None:
|
||||
missing_episodes = [episode for episode in total_episodes if episode not in existing_episodes]
|
||||
total_episode_count = len(total_episodes)
|
||||
else:
|
||||
missing_episodes = None
|
||||
total_episode_count = None
|
||||
seasons_result[str(season)] = {
|
||||
"existing_episodes": existing_episodes,
|
||||
"total_episodes": total_episode_count,
|
||||
"missing_episodes": missing_episodes
|
||||
}
|
||||
if total_episodes is not None and not existing_episodes:
|
||||
missing_seasons.append(season)
|
||||
|
||||
return {
|
||||
"seasons": seasons_result,
|
||||
"missing_seasons": missing_seasons
|
||||
}
|
||||
|
||||
|
||||
class QueryLibraryExistsInput(BaseModel):
|
||||
"""查询媒体库工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
media_type: Optional[str] = Field("all",
|
||||
description="Type of media content: '电影' for films, '电视剧' for television series or anime series, 'all' for all types")
|
||||
title: Optional[str] = Field(None,
|
||||
description="Specific media title to check if it exists in the media library (optional, if provided checks for that specific media)")
|
||||
year: Optional[str] = Field(None,
|
||||
description="Release year of the media (optional, helps narrow down search results)")
|
||||
tmdb_id: Optional[int] = Field(None, description="TMDB ID (can be obtained from search_media tool). Either tmdb_id or douban_id must be provided.")
|
||||
douban_id: Optional[str] = Field(None, description="Douban ID (can be obtained from search_media tool). Either tmdb_id or douban_id must be provided.")
|
||||
media_type: Optional[str] = Field(None, description="Allowed values: movie, tv")
|
||||
|
||||
|
||||
class QueryLibraryExistsTool(MoviePilotTool):
|
||||
name: str = "query_library_exists"
|
||||
description: str = "Check if a specific media resource already exists in the media library (Plex, Emby, Jellyfin). Use this tool to verify whether a movie or TV series has been successfully processed and added to the media server before performing operations like downloading or subscribing."
|
||||
description: str = "Check whether media already exists in Plex, Emby, or Jellyfin by media ID. Results are grouped by media server; TV results include existing episodes, total episodes, and missing episodes/seasons. Requires tmdb_id or douban_id from search_media."
|
||||
args_schema: Type[BaseModel] = QueryLibraryExistsInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据查询参数生成友好的提示消息"""
|
||||
media_type = kwargs.get("media_type", "all")
|
||||
title = kwargs.get("title")
|
||||
year = kwargs.get("year")
|
||||
|
||||
parts = ["正在查询媒体库"]
|
||||
|
||||
if title:
|
||||
parts.append(f"标题: {title}")
|
||||
if year:
|
||||
parts.append(f"年份: {year}")
|
||||
if media_type != "all":
|
||||
parts.append(f"类型: {media_type}")
|
||||
|
||||
return " | ".join(parts) if len(parts) > 1 else parts[0]
|
||||
tmdb_id = kwargs.get("tmdb_id")
|
||||
douban_id = kwargs.get("douban_id")
|
||||
media_type = kwargs.get("media_type")
|
||||
|
||||
async def run(self, media_type: Optional[str] = "all",
|
||||
title: Optional[str] = None, year: Optional[str] = None, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: media_type={media_type}, title={title}")
|
||||
if tmdb_id:
|
||||
message = f"正在查询媒体库: TMDB={tmdb_id}"
|
||||
elif douban_id:
|
||||
message = f"正在查询媒体库: 豆瓣={douban_id}"
|
||||
else:
|
||||
message = "正在查询媒体库"
|
||||
if media_type:
|
||||
message += f" [{media_type}]"
|
||||
return message
|
||||
|
||||
async def run(self, tmdb_id: Optional[int] = None, douban_id: Optional[str] = None,
|
||||
media_type: Optional[str] = None, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: tmdb_id={tmdb_id}, douban_id={douban_id}, media_type={media_type}")
|
||||
try:
|
||||
if not title:
|
||||
return "请提供媒体标题进行查询"
|
||||
|
||||
# 创建 MediaInfo 对象
|
||||
mediainfo = MediaInfo()
|
||||
mediainfo.title = title
|
||||
mediainfo.year = year
|
||||
|
||||
# 转换媒体类型
|
||||
if media_type == "电影":
|
||||
mediainfo.type = MediaType.MOVIE
|
||||
elif media_type == "电视剧":
|
||||
mediainfo.type = MediaType.TV
|
||||
# media_type == "all" 时不设置类型,让媒体服务器自动判断
|
||||
|
||||
# 调用媒体服务器接口实时查询
|
||||
if not tmdb_id and not douban_id:
|
||||
return "参数错误:tmdb_id 和 douban_id 至少需要提供一个,请先使用 search_media 工具获取媒体 ID。"
|
||||
|
||||
media_type_enum = None
|
||||
if media_type:
|
||||
media_type_enum = MediaType.from_agent(media_type)
|
||||
if not media_type_enum:
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv'"
|
||||
|
||||
media_chain = MediaServerChain()
|
||||
existsinfo = media_chain.media_exists(mediainfo=mediainfo)
|
||||
|
||||
if not existsinfo:
|
||||
mediainfo = media_chain.recognize_media(
|
||||
tmdbid=tmdb_id,
|
||||
doubanid=douban_id,
|
||||
mtype=media_type_enum,
|
||||
)
|
||||
if not mediainfo:
|
||||
media_id = f"TMDB={tmdb_id}" if tmdb_id else f"豆瓣={douban_id}"
|
||||
return f"未识别到媒体信息: {media_id}"
|
||||
|
||||
# 2. 遍历所有媒体服务器,分别查询存在性信息
|
||||
server_results = OrderedDict()
|
||||
media_server_helper = MediaServerHelper()
|
||||
total_seasons = _filter_regular_seasons(mediainfo.seasons)
|
||||
global_existsinfo = media_chain.media_exists(mediainfo=mediainfo)
|
||||
|
||||
for service_name in sorted(media_server_helper.get_services().keys()):
|
||||
existsinfo = media_chain.media_exists(mediainfo=mediainfo, server=service_name)
|
||||
if not existsinfo:
|
||||
continue
|
||||
|
||||
if existsinfo.type == MediaType.TV:
|
||||
existing_seasons = _filter_regular_seasons(existsinfo.seasons)
|
||||
server_results[service_name] = _build_tv_server_result(
|
||||
existing_seasons=existing_seasons,
|
||||
total_seasons=total_seasons
|
||||
)
|
||||
else:
|
||||
server_results[service_name] = {
|
||||
"exists": True
|
||||
}
|
||||
|
||||
if global_existsinfo:
|
||||
fallback_server_name = global_existsinfo.server or "local"
|
||||
if fallback_server_name not in server_results:
|
||||
if global_existsinfo.type == MediaType.TV:
|
||||
server_results[fallback_server_name] = _build_tv_server_result(
|
||||
existing_seasons=_filter_regular_seasons(global_existsinfo.seasons),
|
||||
total_seasons=total_seasons
|
||||
)
|
||||
else:
|
||||
server_results[fallback_server_name] = {
|
||||
"exists": True
|
||||
}
|
||||
|
||||
if not server_results:
|
||||
return "媒体库中未找到相关媒体"
|
||||
|
||||
# 如果找到了,获取详细信息
|
||||
result_items = []
|
||||
if existsinfo.itemid and existsinfo.server:
|
||||
iteminfo = media_chain.iteminfo(server=existsinfo.server, item_id=existsinfo.itemid)
|
||||
if iteminfo:
|
||||
# 使用 model_dump() 转换为字典格式
|
||||
item_dict = iteminfo.model_dump(exclude_none=True)
|
||||
result_items.append(item_dict)
|
||||
|
||||
if result_items:
|
||||
return json.dumps(result_items, ensure_ascii=False)
|
||||
|
||||
# 如果找到了但没有详细信息,返回基本信息
|
||||
|
||||
# 3. 组装统一的存在性结果,不查询媒体服务器详情
|
||||
result_dict = {
|
||||
"type": existsinfo.type.value if existsinfo.type else None,
|
||||
"server": existsinfo.server,
|
||||
"server_type": existsinfo.server_type,
|
||||
"itemid": existsinfo.itemid,
|
||||
"seasons": existsinfo.seasons if existsinfo.seasons else {}
|
||||
"title": mediainfo.title,
|
||||
"year": mediainfo.year,
|
||||
"type": media_type_to_agent(mediainfo.type),
|
||||
"servers": server_results
|
||||
}
|
||||
|
||||
return json.dumps([result_dict], ensure_ascii=False)
|
||||
except Exception as e:
|
||||
logger.error(f"查询媒体库失败: {e}", exc_info=True)
|
||||
|
||||
126
app/agent/tools/impl/query_media_detail.py
Normal file
126
app/agent/tools/impl/query_media_detail.py
Normal file
@@ -0,0 +1,126 @@
|
||||
"""查询媒体详情工具"""
|
||||
|
||||
import json
|
||||
from typing import Optional, Type
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.media import MediaChain
|
||||
from app.log import logger
|
||||
from app.schemas.types import MediaType
|
||||
|
||||
|
||||
class QueryMediaDetailInput(BaseModel):
|
||||
"""查询媒体详情工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
tmdb_id: Optional[int] = Field(None, description="TMDB ID of the media (movie or TV series, can be obtained from search_media tool)")
|
||||
douban_id: Optional[str] = Field(None, description="Douban ID of the media (alternative to tmdb_id)")
|
||||
media_type: str = Field(..., description="Allowed values: movie, tv")
|
||||
|
||||
|
||||
class QueryMediaDetailTool(MoviePilotTool):
|
||||
name: str = "query_media_detail"
|
||||
description: str = "Query supplementary media details from TMDB by ID and media_type. Accepts tmdb_id or douban_id (at least one required). media_type accepts 'movie' or 'tv'. Returns non-duplicated detail fields such as status, genres, directors, actors, and season info for TV series."
|
||||
args_schema: Type[BaseModel] = QueryMediaDetailInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据查询参数生成友好的提示消息"""
|
||||
tmdb_id = kwargs.get("tmdb_id")
|
||||
douban_id = kwargs.get("douban_id")
|
||||
if tmdb_id:
|
||||
return f"正在查询媒体详情: TMDB ID {tmdb_id}"
|
||||
return f"正在查询媒体详情: 豆瓣 ID {douban_id}"
|
||||
|
||||
async def run(self, media_type: str, tmdb_id: Optional[int] = None, douban_id: Optional[str] = None, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: tmdb_id={tmdb_id}, douban_id={douban_id}, media_type={media_type}")
|
||||
|
||||
if tmdb_id is None and douban_id is None:
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"message": "必须提供 tmdb_id 或 douban_id 之一"
|
||||
}, ensure_ascii=False)
|
||||
|
||||
try:
|
||||
media_chain = MediaChain()
|
||||
|
||||
media_type_enum = MediaType.from_agent(media_type)
|
||||
if not media_type_enum:
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"message": f"无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv'"
|
||||
}, ensure_ascii=False)
|
||||
|
||||
mediainfo = await media_chain.async_recognize_media(tmdbid=tmdb_id, doubanid=douban_id, mtype=media_type_enum)
|
||||
|
||||
if not mediainfo:
|
||||
id_info = f"TMDB ID {tmdb_id}" if tmdb_id else f"豆瓣 ID {douban_id}"
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"message": f"未找到 {id_info} 的媒体信息"
|
||||
}, ensure_ascii=False)
|
||||
|
||||
# 精简 genres - 只保留名称
|
||||
genres = [g.get("name") for g in (mediainfo.genres or []) if g.get("name")]
|
||||
|
||||
# 精简 directors - 只保留姓名和职位
|
||||
directors = [
|
||||
{
|
||||
"name": d.get("name"),
|
||||
"job": d.get("job")
|
||||
}
|
||||
for d in (mediainfo.directors or [])
|
||||
if d.get("name")
|
||||
]
|
||||
|
||||
# 精简 actors - 只保留姓名和角色
|
||||
actors = [
|
||||
{
|
||||
"name": a.get("name"),
|
||||
"character": a.get("character")
|
||||
}
|
||||
for a in (mediainfo.actors or [])
|
||||
if a.get("name")
|
||||
]
|
||||
|
||||
# 构建基础媒体详情信息
|
||||
result = {
|
||||
"status": mediainfo.status,
|
||||
"genres": genres,
|
||||
"directors": directors,
|
||||
"actors": actors
|
||||
}
|
||||
|
||||
# 如果是电视剧,添加电视剧特有信息
|
||||
if mediainfo.type == MediaType.TV:
|
||||
# 精简 season_info - 只保留基础摘要
|
||||
season_info = [
|
||||
{
|
||||
"season_number": s.get("season_number"),
|
||||
"name": s.get("name"),
|
||||
"episode_count": s.get("episode_count"),
|
||||
"air_date": s.get("air_date")
|
||||
}
|
||||
for s in (mediainfo.season_info or [])
|
||||
if s.get("season_number") is not None
|
||||
]
|
||||
|
||||
result.update({
|
||||
"number_of_seasons": mediainfo.number_of_seasons,
|
||||
"number_of_episodes": mediainfo.number_of_episodes,
|
||||
"first_air_date": mediainfo.first_air_date,
|
||||
"last_air_date": mediainfo.last_air_date,
|
||||
"season_info": season_info
|
||||
})
|
||||
|
||||
return json.dumps(result, ensure_ascii=False, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
error_message = f"查询媒体详情失败: {str(e)}"
|
||||
logger.error(f"查询媒体详情失败: {e}", exc_info=True)
|
||||
return json.dumps({
|
||||
"success": False,
|
||||
"message": error_message,
|
||||
"tmdb_id": tmdb_id,
|
||||
"douban_id": douban_id
|
||||
}, ensure_ascii=False)
|
||||
@@ -10,13 +10,13 @@ from app.agent.tools.base import MoviePilotTool
|
||||
from app.core.context import MediaInfo
|
||||
from app.helper.subscribe import SubscribeHelper
|
||||
from app.log import logger
|
||||
from app.schemas.types import MediaType
|
||||
from app.schemas.types import MediaType, media_type_to_agent
|
||||
|
||||
|
||||
class QueryPopularSubscribesInput(BaseModel):
|
||||
"""查询热门订阅工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
stype: str = Field(..., description="Media type: '电影' for films, '电视剧' for television series")
|
||||
media_type: str = Field(..., description="Allowed values: movie, tv")
|
||||
page: Optional[int] = Field(1, description="Page number for pagination (default: 1)")
|
||||
count: Optional[int] = Field(30, description="Number of items per page (default: 30)")
|
||||
min_sub: Optional[int] = Field(None, description="Minimum number of subscribers filter (optional, e.g., 5)")
|
||||
@@ -33,13 +33,13 @@ class QueryPopularSubscribesTool(MoviePilotTool):
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据查询参数生成友好的提示消息"""
|
||||
stype = kwargs.get("stype", "")
|
||||
media_type = kwargs.get("media_type", "")
|
||||
page = kwargs.get("page", 1)
|
||||
min_sub = kwargs.get("min_sub")
|
||||
min_rating = kwargs.get("min_rating")
|
||||
max_rating = kwargs.get("max_rating")
|
||||
|
||||
parts = [f"正在查询热门订阅 [{stype}]"]
|
||||
parts = [f"正在查询热门订阅 [{media_type}]"]
|
||||
|
||||
if min_sub:
|
||||
parts.append(f"最少订阅: {min_sub}")
|
||||
@@ -52,7 +52,7 @@ class QueryPopularSubscribesTool(MoviePilotTool):
|
||||
|
||||
return " | ".join(parts) if len(parts) > 1 else parts[0]
|
||||
|
||||
async def run(self, stype: str,
|
||||
async def run(self, media_type: str,
|
||||
page: Optional[int] = 1,
|
||||
count: Optional[int] = 30,
|
||||
min_sub: Optional[int] = None,
|
||||
@@ -61,7 +61,7 @@ class QueryPopularSubscribesTool(MoviePilotTool):
|
||||
max_rating: Optional[float] = None,
|
||||
sort_type: Optional[str] = None, **kwargs) -> str:
|
||||
logger.info(
|
||||
f"执行工具: {self.name}, 参数: stype={stype}, page={page}, count={count}, min_sub={min_sub}, "
|
||||
f"执行工具: {self.name}, 参数: media_type={media_type}, page={page}, count={count}, min_sub={min_sub}, "
|
||||
f"genre_id={genre_id}, min_rating={min_rating}, max_rating={max_rating}, sort_type={sort_type}")
|
||||
|
||||
try:
|
||||
@@ -69,10 +69,13 @@ class QueryPopularSubscribesTool(MoviePilotTool):
|
||||
page = 1
|
||||
if count is None or count < 1:
|
||||
count = 30
|
||||
media_type_enum = MediaType.from_agent(media_type)
|
||||
if not media_type_enum:
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv'"
|
||||
|
||||
subscribe_helper = SubscribeHelper()
|
||||
subscribes = await subscribe_helper.async_get_statistic(
|
||||
stype=stype,
|
||||
stype=media_type_enum.to_agent(),
|
||||
page=page,
|
||||
count=count,
|
||||
genre_id=genre_id,
|
||||
@@ -94,7 +97,15 @@ class QueryPopularSubscribesTool(MoviePilotTool):
|
||||
continue
|
||||
|
||||
media = MediaInfo()
|
||||
media.type = MediaType(sub.get("type"))
|
||||
raw_type = str(sub.get("type") or "").strip().lower()
|
||||
if raw_type in ["movie", "电影"]:
|
||||
media.type = MediaType.MOVIE
|
||||
elif raw_type in ["tv", "电视剧"]:
|
||||
media.type = MediaType.TV
|
||||
else:
|
||||
# 跳过无法识别类型的数据,避免单条脏数据导致整批失败
|
||||
logger.warning(f"跳过未知媒体类型: {sub.get('type')}")
|
||||
continue
|
||||
media.tmdb_id = sub.get("tmdbid")
|
||||
# 处理标题
|
||||
title = sub.get("name")
|
||||
@@ -124,7 +135,7 @@ class QueryPopularSubscribesTool(MoviePilotTool):
|
||||
for media in ret_medias:
|
||||
media_dict = media.to_dict()
|
||||
simplified = {
|
||||
"type": media_dict.get("type"),
|
||||
"type": media_type_to_agent(media_dict.get("type")),
|
||||
"title": media_dict.get("title"),
|
||||
"year": media_dict.get("year"),
|
||||
"tmdb_id": media_dict.get("tmdb_id"),
|
||||
|
||||
@@ -15,7 +15,7 @@ from app.log import logger
|
||||
class QuerySiteUserdataInput(BaseModel):
|
||||
"""查询站点用户数据工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
site_id: int = Field(..., description="The ID of the site to query user data for")
|
||||
site_id: int = Field(..., description="The ID of the site to query user data for (can be obtained from query_sites tool)")
|
||||
workdate: Optional[str] = Field(None, description="Work date to query (optional, format: 'YYYY-MM-DD', if not specified returns latest data)")
|
||||
|
||||
|
||||
|
||||
@@ -9,12 +9,13 @@ from app.agent.tools.base import MoviePilotTool
|
||||
from app.db import AsyncSessionFactory
|
||||
from app.db.models.subscribehistory import SubscribeHistory
|
||||
from app.log import logger
|
||||
from app.schemas.types import media_type_to_agent
|
||||
|
||||
|
||||
class QuerySubscribeHistoryInput(BaseModel):
|
||||
"""查询订阅历史工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
media_type: Optional[str] = Field("all", description="Filter by media type: '电影' for films, '电视剧' for television series, 'all' for all types (default: 'all')")
|
||||
media_type: Optional[str] = Field("all", description="Allowed values: movie, tv, all")
|
||||
name: Optional[str] = Field(None, description="Filter by media name (partial match, optional)")
|
||||
|
||||
|
||||
@@ -42,6 +43,9 @@ class QuerySubscribeHistoryTool(MoviePilotTool):
|
||||
logger.info(f"执行工具: {self.name}, 参数: media_type={media_type}, name={name}")
|
||||
|
||||
try:
|
||||
if media_type not in ["all", "movie", "tv"]:
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv', 'all'"
|
||||
|
||||
# 获取数据库会话
|
||||
async with AsyncSessionFactory() as db:
|
||||
# 根据类型查询
|
||||
@@ -80,7 +84,7 @@ class QuerySubscribeHistoryTool(MoviePilotTool):
|
||||
"id": record.id,
|
||||
"name": record.name,
|
||||
"year": record.year,
|
||||
"type": record.type,
|
||||
"type": media_type_to_agent(record.type),
|
||||
"season": record.season,
|
||||
"tmdbid": record.tmdbid,
|
||||
"doubanid": record.doubanid,
|
||||
|
||||
@@ -8,20 +8,51 @@ from pydantic import BaseModel, Field
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.db.subscribe_oper import SubscribeOper
|
||||
from app.log import logger
|
||||
from app.schemas.subscribe import Subscribe as SubscribeSchema
|
||||
from app.schemas.types import MediaType
|
||||
|
||||
QUERY_SUBSCRIBE_OUTPUT_FIELDS = [
|
||||
"id",
|
||||
"name",
|
||||
"year",
|
||||
"type",
|
||||
"season",
|
||||
"total_episode",
|
||||
"start_episode",
|
||||
"lack_episode",
|
||||
"filter",
|
||||
"include",
|
||||
"exclude",
|
||||
"quality",
|
||||
"resolution",
|
||||
"effect",
|
||||
"state",
|
||||
"last_update",
|
||||
"sites",
|
||||
"downloader",
|
||||
"best_version",
|
||||
"save_path",
|
||||
"custom_words",
|
||||
"media_category",
|
||||
"filter_groups",
|
||||
"episode_group"
|
||||
]
|
||||
|
||||
|
||||
class QuerySubscribesInput(BaseModel):
|
||||
"""查询订阅工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
status: Optional[str] = Field("all",
|
||||
description="Filter subscriptions by status: 'R' for enabled subscriptions, 'P' for disabled ones, 'all' for all subscriptions")
|
||||
description="Filter subscriptions by status: 'R' for enabled subscriptions, 'S' for paused ones, 'all' for all subscriptions")
|
||||
media_type: Optional[str] = Field("all",
|
||||
description="Filter by media type: '电影' for films, '电视剧' for television series, 'all' for all types")
|
||||
description="Allowed values: movie, tv, all")
|
||||
tmdb_id: Optional[int] = Field(None, description="Filter by TMDB ID to check if a specific media is already subscribed")
|
||||
douban_id: Optional[str] = Field(None, description="Filter by Douban ID to check if a specific media is already subscribed")
|
||||
|
||||
|
||||
class QuerySubscribesTool(MoviePilotTool):
|
||||
name: str = "query_subscribes"
|
||||
description: str = "Query subscription status and list all user subscriptions. Shows active subscriptions, their download status, and configuration details."
|
||||
description: str = "Query subscription status and list user subscriptions. Returns full subscription parameters for each matched subscription."
|
||||
args_schema: Type[BaseModel] = QuerySubscribesInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
@@ -33,7 +64,7 @@ class QuerySubscribesTool(MoviePilotTool):
|
||||
|
||||
# 根据状态过滤条件生成提示
|
||||
if status != "all":
|
||||
status_map = {"R": "已启用", "P": "已禁用"}
|
||||
status_map = {"R": "已启用", "S": "已暂停"}
|
||||
parts.append(f"状态: {status_map.get(status, status)}")
|
||||
|
||||
# 根据媒体类型过滤条件生成提示
|
||||
@@ -42,44 +73,38 @@ class QuerySubscribesTool(MoviePilotTool):
|
||||
|
||||
return " | ".join(parts) if len(parts) > 1 else parts[0]
|
||||
|
||||
async def run(self, status: Optional[str] = "all", media_type: Optional[str] = "all", **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: status={status}, media_type={media_type}")
|
||||
async def run(self, status: Optional[str] = "all", media_type: Optional[str] = "all",
|
||||
tmdb_id: Optional[int] = None, douban_id: Optional[str] = None, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: status={status}, media_type={media_type}, tmdb_id={tmdb_id}, douban_id={douban_id}")
|
||||
try:
|
||||
if media_type != "all" and not MediaType.from_agent(media_type):
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv', 'all'"
|
||||
|
||||
subscribe_oper = SubscribeOper()
|
||||
subscribes = await subscribe_oper.async_list()
|
||||
filtered_subscribes = []
|
||||
for sub in subscribes:
|
||||
if status != "all" and sub.state != status:
|
||||
continue
|
||||
if media_type != "all" and sub.type != media_type:
|
||||
if media_type != "all" and sub.type != MediaType.from_agent(media_type).value:
|
||||
continue
|
||||
if tmdb_id is not None and sub.tmdbid != tmdb_id:
|
||||
continue
|
||||
if douban_id is not None and sub.doubanid != douban_id:
|
||||
continue
|
||||
filtered_subscribes.append(sub)
|
||||
if filtered_subscribes:
|
||||
# 限制最多50条结果
|
||||
total_count = len(filtered_subscribes)
|
||||
limited_subscribes = filtered_subscribes[:50]
|
||||
# 精简字段,只保留关键信息
|
||||
simplified_subscribes = []
|
||||
for s in limited_subscribes:
|
||||
simplified = {
|
||||
"id": s.id,
|
||||
"name": s.name,
|
||||
"year": s.year,
|
||||
"type": s.type,
|
||||
"season": s.season,
|
||||
"tmdbid": s.tmdbid,
|
||||
"doubanid": s.doubanid,
|
||||
"bangumiid": s.bangumiid,
|
||||
"poster": s.poster,
|
||||
"vote": s.vote,
|
||||
"state": s.state,
|
||||
"total_episode": s.total_episode,
|
||||
"lack_episode": s.lack_episode,
|
||||
"last_update": s.last_update,
|
||||
"username": s.username
|
||||
}
|
||||
simplified_subscribes.append(simplified)
|
||||
result_json = json.dumps(simplified_subscribes, ensure_ascii=False, indent=2)
|
||||
full_subscribes = [
|
||||
SubscribeSchema.model_validate(s, from_attributes=True).model_dump(
|
||||
include=set(QUERY_SUBSCRIBE_OUTPUT_FIELDS),
|
||||
exclude_none=True
|
||||
)
|
||||
for s in limited_subscribes
|
||||
]
|
||||
result_json = json.dumps(full_subscribes, ensure_ascii=False, indent=2)
|
||||
# 如果结果被裁剪,添加提示信息
|
||||
if total_count > 50:
|
||||
return f"注意:查询结果共找到 {total_count} 条,为节省上下文空间,仅显示前 50 条结果。\n\n{result_json}"
|
||||
|
||||
@@ -10,6 +10,7 @@ from app.agent.tools.base import MoviePilotTool
|
||||
from app.db import AsyncSessionFactory
|
||||
from app.db.models.transferhistory import TransferHistory
|
||||
from app.log import logger
|
||||
from app.schemas.types import media_type_to_agent
|
||||
|
||||
|
||||
class QueryTransferHistoryInput(BaseModel):
|
||||
@@ -95,7 +96,7 @@ class QueryTransferHistoryTool(MoviePilotTool):
|
||||
"id": record.id,
|
||||
"title": record.title,
|
||||
"year": record.year,
|
||||
"type": record.type,
|
||||
"type": media_type_to_agent(record.type),
|
||||
"category": record.category,
|
||||
"seasons": record.seasons,
|
||||
"episodes": record.episodes,
|
||||
|
||||
81
app/agent/tools/impl/read_file.py
Normal file
81
app/agent/tools/impl/read_file.py
Normal file
@@ -0,0 +1,81 @@
|
||||
"""文件读取工具"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Optional, Type
|
||||
|
||||
from anyio import Path as AsyncPath
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.log import logger
|
||||
|
||||
# 最大读取大小 50KB
|
||||
MAX_READ_SIZE = 50 * 1024
|
||||
|
||||
|
||||
class ReadFileInput(BaseModel):
|
||||
"""Input parameters for read file tool"""
|
||||
file_path: str = Field(..., description="The absolute path of the file to read")
|
||||
start_line: Optional[int] = Field(None, description="The starting line number (1-based, inclusive). If not provided, reading starts from the beginning of the file.")
|
||||
end_line: Optional[int] = Field(None, description="The ending line number (1-based, inclusive). If not provided, reading goes until the end of the file.")
|
||||
|
||||
|
||||
class ReadFileTool(MoviePilotTool):
|
||||
name: str = "read_file"
|
||||
description: str = "Read the content of a text file. Supports reading by line range. Each read is limited to 50KB; content exceeding this limit will be truncated."
|
||||
args_schema: Type[BaseModel] = ReadFileInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据参数生成友好的提示消息"""
|
||||
file_path = kwargs.get("file_path", "")
|
||||
file_name = Path(file_path).name if file_path else "未知文件"
|
||||
return f"正在读取文件: {file_name}"
|
||||
|
||||
async def run(self, file_path: str, start_line: Optional[int] = None,
|
||||
end_line: Optional[int] = None, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: file_path={file_path}, start_line={start_line}, end_line={end_line}")
|
||||
|
||||
try:
|
||||
path = AsyncPath(file_path)
|
||||
|
||||
if not await path.exists():
|
||||
return f"错误:文件 {file_path} 不存在"
|
||||
|
||||
if not await path.is_file():
|
||||
return f"错误:{file_path} 不是一个文件"
|
||||
|
||||
content = await path.read_text(encoding="utf-8")
|
||||
truncated = False
|
||||
|
||||
if start_line is not None or end_line is not None:
|
||||
lines = content.splitlines(keepends=True)
|
||||
total_lines = len(lines)
|
||||
|
||||
# 将行号转换为索引(1-based -> 0-based)
|
||||
s = (start_line - 1) if start_line and start_line >= 1 else 0
|
||||
e = end_line if end_line and end_line >= 1 else total_lines
|
||||
|
||||
# 确保范围有效
|
||||
s = max(0, min(s, total_lines))
|
||||
e = max(s, min(e, total_lines))
|
||||
|
||||
content = "".join(lines[s:e])
|
||||
|
||||
# 检查大小限制
|
||||
content_bytes = content.encode("utf-8")
|
||||
if len(content_bytes) > MAX_READ_SIZE:
|
||||
content = content_bytes[:MAX_READ_SIZE].decode("utf-8", errors="ignore")
|
||||
truncated = True
|
||||
|
||||
if truncated:
|
||||
return f"{content}\n\n[警告:文件内容已超过50KB限制,以上内容已被截断。请使用 start_line/end_line 参数分段读取。]"
|
||||
|
||||
return content
|
||||
|
||||
except PermissionError:
|
||||
return f"错误:没有权限读取 {file_path}"
|
||||
except UnicodeDecodeError:
|
||||
return f"错误:{file_path} 不是文本文件,无法读取"
|
||||
except Exception as e:
|
||||
logger.error(f"读取文件 {file_path} 时发生错误: {str(e)}", exc_info=True)
|
||||
return f"操作失败: {str(e)}"
|
||||
@@ -10,6 +10,7 @@ from app.chain.media import MediaChain
|
||||
from app.core.context import Context
|
||||
from app.core.metainfo import MetaInfo
|
||||
from app.log import logger
|
||||
from app.schemas.types import media_type_to_agent
|
||||
|
||||
|
||||
class RecognizeMediaInput(BaseModel):
|
||||
@@ -124,7 +125,7 @@ class RecognizeMediaTool(MoviePilotTool):
|
||||
"title": media_info.get("title"),
|
||||
"en_title": media_info.get("en_title"),
|
||||
"year": media_info.get("year"),
|
||||
"type": media_info.get("type"),
|
||||
"type": media_type_to_agent(media_info.get("type")),
|
||||
"season": media_info.get("season"),
|
||||
"tmdb_id": media_info.get("tmdb_id"),
|
||||
"imdb_id": media_info.get("imdb_id"),
|
||||
@@ -145,7 +146,7 @@ class RecognizeMediaTool(MoviePilotTool):
|
||||
"name": meta_info.get("name"),
|
||||
"title": meta_info.get("title"),
|
||||
"year": meta_info.get("year"),
|
||||
"type": meta_info.get("type"),
|
||||
"type": media_type_to_agent(meta_info.get("type")),
|
||||
"begin_season": meta_info.get("begin_season"),
|
||||
"end_season": meta_info.get("end_season"),
|
||||
"begin_episode": meta_info.get("begin_episode"),
|
||||
|
||||
@@ -14,21 +14,21 @@ from app.log import logger
|
||||
class RunWorkflowInput(BaseModel):
|
||||
"""执行工作流工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
workflow_identifier: str = Field(..., description="Workflow identifier: can be workflow ID (integer as string) or workflow name")
|
||||
workflow_id: int = Field(..., description="Workflow ID (can be obtained from query_workflows tool)")
|
||||
from_begin: Optional[bool] = Field(True, description="Whether to run workflow from the beginning (default: True, if False will continue from last executed action)")
|
||||
|
||||
|
||||
class RunWorkflowTool(MoviePilotTool):
|
||||
name: str = "run_workflow"
|
||||
description: str = "Execute a specific workflow manually. Can run workflow by ID or name. Supports running from the beginning or continuing from the last executed action."
|
||||
description: str = "Execute a specific workflow manually by workflow ID. Supports running from the beginning or continuing from the last executed action."
|
||||
args_schema: Type[BaseModel] = RunWorkflowInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据工作流参数生成友好的提示消息"""
|
||||
workflow_identifier = kwargs.get("workflow_identifier", "")
|
||||
workflow_id = kwargs.get("workflow_id")
|
||||
from_begin = kwargs.get("from_begin", True)
|
||||
|
||||
message = f"正在执行工作流: {workflow_identifier}"
|
||||
message = f"正在执行工作流: {workflow_id}"
|
||||
if not from_begin:
|
||||
message += " (从上次位置继续)"
|
||||
else:
|
||||
@@ -36,27 +36,18 @@ class RunWorkflowTool(MoviePilotTool):
|
||||
|
||||
return message
|
||||
|
||||
async def run(self, workflow_identifier: str,
|
||||
async def run(self, workflow_id: int,
|
||||
from_begin: Optional[bool] = True, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: workflow_identifier={workflow_identifier}, from_begin={from_begin}")
|
||||
logger.info(f"执行工具: {self.name}, 参数: workflow_id={workflow_id}, from_begin={from_begin}")
|
||||
|
||||
try:
|
||||
# 获取数据库会话
|
||||
async with AsyncSessionFactory() as db:
|
||||
workflow_oper = WorkflowOper(db)
|
||||
|
||||
# 尝试解析为工作流ID
|
||||
workflow = None
|
||||
if workflow_identifier.isdigit():
|
||||
# 如果是数字,尝试作为工作流ID查询
|
||||
workflow = await workflow_oper.async_get(int(workflow_identifier))
|
||||
|
||||
# 如果不是ID或ID查询失败,尝试按名称查询
|
||||
if not workflow:
|
||||
workflow = await workflow_oper.async_get_by_name(workflow_identifier)
|
||||
workflow = await workflow_oper.async_get(workflow_id)
|
||||
|
||||
if not workflow:
|
||||
return f"未找到工作流:{workflow_identifier},请使用 query_workflows 工具查询可用的工作流"
|
||||
return f"未找到工作流:{workflow_id},请使用 query_workflows 工具查询可用的工作流"
|
||||
|
||||
# 执行工作流
|
||||
workflow_chain = WorkflowChain()
|
||||
|
||||
@@ -8,7 +8,7 @@ from pydantic import BaseModel, Field
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.media import MediaChain
|
||||
from app.log import logger
|
||||
from app.schemas.types import MediaType
|
||||
from app.schemas.types import MediaType, media_type_to_agent
|
||||
|
||||
|
||||
class SearchMediaInput(BaseModel):
|
||||
@@ -17,7 +17,7 @@ class SearchMediaInput(BaseModel):
|
||||
title: str = Field(..., description="The title of the media to search for (e.g., 'The Matrix', 'Breaking Bad')")
|
||||
year: Optional[str] = Field(None, description="Release year of the media (optional, helps narrow down results)")
|
||||
media_type: Optional[str] = Field(None,
|
||||
description="Type of media content: '电影' for films, '电视剧' for television series or anime series")
|
||||
description="Allowed values: movie, tv")
|
||||
season: Optional[int] = Field(None,
|
||||
description="Season number for TV shows and anime (optional, only applicable for series)")
|
||||
|
||||
@@ -56,14 +56,19 @@ class SearchMediaTool(MoviePilotTool):
|
||||
|
||||
# 过滤结果
|
||||
if results:
|
||||
media_type_enum = None
|
||||
if media_type:
|
||||
media_type_enum = MediaType.from_agent(media_type)
|
||||
if not media_type_enum:
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv'"
|
||||
|
||||
filtered_results = []
|
||||
for result in results:
|
||||
if year and result.year != year:
|
||||
continue
|
||||
if media_type:
|
||||
if result.type != MediaType(media_type):
|
||||
continue
|
||||
if season and result.season != season:
|
||||
if media_type_enum and result.type != media_type_enum:
|
||||
continue
|
||||
if season is not None and result.season != season:
|
||||
continue
|
||||
filtered_results.append(result)
|
||||
|
||||
@@ -78,7 +83,7 @@ class SearchMediaTool(MoviePilotTool):
|
||||
"title": r.title,
|
||||
"en_title": r.en_title,
|
||||
"year": r.year,
|
||||
"type": r.type.value if r.type else None,
|
||||
"type": media_type_to_agent(r.type),
|
||||
"season": r.season,
|
||||
"tmdb_id": r.tmdb_id,
|
||||
"imdb_id": r.imdb_id,
|
||||
|
||||
@@ -10,15 +10,16 @@ from app.chain.subscribe import SubscribeChain
|
||||
from app.core.config import global_vars
|
||||
from app.db.subscribe_oper import SubscribeOper
|
||||
from app.log import logger
|
||||
from app.schemas.types import media_type_to_agent
|
||||
|
||||
|
||||
class SearchSubscribeInput(BaseModel):
|
||||
"""搜索订阅缺失剧集工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
subscribe_id: int = Field(..., description="The ID of the subscription to search for missing episodes")
|
||||
subscribe_id: int = Field(..., description="The ID of the subscription to search for missing episodes (can be obtained from query_subscribes tool)")
|
||||
manual: Optional[bool] = Field(False, description="Whether this is a manual search (default: False)")
|
||||
filter_groups: Optional[List[str]] = Field(None,
|
||||
description="List of filter rule group names to apply for this search (optional, use query_rule_groups tool to get available rule groups. If provided, will temporarily update the subscription's filter groups before searching)")
|
||||
description="List of filter rule group names to apply for this search (optional, can be obtained from query_rule_groups tool. If provided, will temporarily update the subscription's filter groups before searching)")
|
||||
|
||||
|
||||
class SearchSubscribeTool(MoviePilotTool):
|
||||
@@ -58,7 +59,7 @@ class SearchSubscribeTool(MoviePilotTool):
|
||||
"id": subscribe.id,
|
||||
"name": subscribe.name,
|
||||
"year": subscribe.year,
|
||||
"type": subscribe.type,
|
||||
"type": media_type_to_agent(subscribe.type),
|
||||
"season": subscribe.season,
|
||||
"state": subscribe.state,
|
||||
"total_episode": subscribe.total_episode,
|
||||
|
||||
@@ -1,141 +1,109 @@
|
||||
"""搜索种子工具"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import List, Optional, Type
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.search import SearchChain
|
||||
from app.db.systemconfig_oper import SystemConfigOper
|
||||
from app.helper.sites import SitesHelper
|
||||
from app.log import logger
|
||||
from app.schemas.types import MediaType
|
||||
from app.schemas.types import MediaType, SystemConfigKey
|
||||
from ._torrent_search_utils import (
|
||||
SEARCH_RESULT_CACHE_FILE,
|
||||
build_filter_options,
|
||||
)
|
||||
|
||||
|
||||
class SearchTorrentsInput(BaseModel):
|
||||
"""搜索种子工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
title: str = Field(...,
|
||||
description="The title of the media resource to search for (e.g., 'The Matrix 1999', 'Breaking Bad S01E01')")
|
||||
year: Optional[str] = Field(None,
|
||||
description="Release year of the media (optional, helps narrow down search results)")
|
||||
media_type: Optional[str] = Field(None,
|
||||
description="Type of media content: '电影' for films, '电视剧' for television series or anime series")
|
||||
season: Optional[int] = Field(None, description="Season number for TV shows (optional, only applicable for series)")
|
||||
tmdb_id: Optional[int] = Field(None, description="TMDB ID (can be obtained from search_media tool). Either tmdb_id or douban_id must be provided.")
|
||||
douban_id: Optional[str] = Field(None, description="Douban ID (can be obtained from search_media tool). Either tmdb_id or douban_id must be provided.")
|
||||
media_type: Optional[str] = Field(None, description="Allowed values: movie, tv")
|
||||
area: Optional[str] = Field(None, description="Search scope: 'title' (default) or 'imdbid'")
|
||||
sites: Optional[List[int]] = Field(None,
|
||||
description="Array of specific site IDs to search on (optional, if not provided searches all configured sites)")
|
||||
filter_pattern: Optional[str] = Field(None,
|
||||
description="Regular expression pattern to filter torrent titles by resolution, quality, or other keywords (e.g., '4K|2160p|UHD' for 4K content, '1080p|BluRay' for 1080p BluRay)")
|
||||
|
||||
|
||||
class SearchTorrentsTool(MoviePilotTool):
|
||||
name: str = "search_torrents"
|
||||
description: str = "Search for torrent files across configured indexer sites based on media information. Returns available torrent downloads with details like file size, quality, and download links."
|
||||
description: str = ("Search for torrent files by media ID across configured indexer sites, cache the matched results, "
|
||||
"and return available filter options for follow-up selection. "
|
||||
"Requires tmdb_id or douban_id (can be obtained from search_media tool) for accurate matching.")
|
||||
args_schema: Type[BaseModel] = SearchTorrentsInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据搜索参数生成友好的提示消息"""
|
||||
title = kwargs.get("title", "")
|
||||
year = kwargs.get("year")
|
||||
tmdb_id = kwargs.get("tmdb_id")
|
||||
douban_id = kwargs.get("douban_id")
|
||||
media_type = kwargs.get("media_type")
|
||||
season = kwargs.get("season")
|
||||
filter_pattern = kwargs.get("filter_pattern")
|
||||
|
||||
message = f"正在搜索种子: {title}"
|
||||
if year:
|
||||
message += f" ({year})"
|
||||
|
||||
if tmdb_id:
|
||||
message = f"正在搜索种子: TMDB={tmdb_id}"
|
||||
elif douban_id:
|
||||
message = f"正在搜索种子: 豆瓣={douban_id}"
|
||||
else:
|
||||
message = "正在搜索种子"
|
||||
if media_type:
|
||||
message += f" [{media_type}]"
|
||||
if season:
|
||||
message += f" 第{season}季"
|
||||
if filter_pattern:
|
||||
message += f" 过滤: {filter_pattern}"
|
||||
|
||||
return message
|
||||
|
||||
async def run(self, title: str, year: Optional[str] = None,
|
||||
media_type: Optional[str] = None, season: Optional[int] = None,
|
||||
sites: Optional[List[int]] = None, filter_pattern: Optional[str] = None, **kwargs) -> str:
|
||||
async def run(self, tmdb_id: Optional[int] = None, douban_id: Optional[str] = None,
|
||||
media_type: Optional[str] = None, area: Optional[str] = None,
|
||||
sites: Optional[List[int]] = None, **kwargs) -> str:
|
||||
logger.info(
|
||||
f"执行工具: {self.name}, 参数: title={title}, year={year}, media_type={media_type}, season={season}, sites={sites}, filter_pattern={filter_pattern}")
|
||||
f"执行工具: {self.name}, 参数: tmdb_id={tmdb_id}, douban_id={douban_id}, media_type={media_type}, area={area}, sites={sites}")
|
||||
|
||||
if not tmdb_id and not douban_id:
|
||||
return "参数错误:tmdb_id 和 douban_id 至少需要提供一个,请先使用 search_media 工具获取媒体 ID。"
|
||||
|
||||
try:
|
||||
search_chain = SearchChain()
|
||||
torrents = await search_chain.async_search_by_title(title=title, sites=sites)
|
||||
filtered_torrents = []
|
||||
# 编译正则表达式(如果提供)
|
||||
regex_pattern = None
|
||||
if filter_pattern:
|
||||
try:
|
||||
regex_pattern = re.compile(filter_pattern, re.IGNORECASE)
|
||||
except re.error as e:
|
||||
logger.warning(f"正则表达式编译失败: {filter_pattern}, 错误: {e}")
|
||||
return f"正则表达式格式错误: {str(e)}"
|
||||
|
||||
for torrent in torrents:
|
||||
# torrent 是 Context 对象,需要通过 meta_info 和 media_info 访问属性
|
||||
if year and torrent.meta_info and torrent.meta_info.year != year:
|
||||
continue
|
||||
if media_type and torrent.media_info:
|
||||
if torrent.media_info.type != MediaType(media_type):
|
||||
continue
|
||||
if season and torrent.meta_info and torrent.meta_info.begin_season != season:
|
||||
continue
|
||||
# 使用正则表达式过滤标题(分辨率、质量等关键字)
|
||||
if regex_pattern and torrent.torrent_info and torrent.torrent_info.title:
|
||||
if not regex_pattern.search(torrent.torrent_info.title):
|
||||
continue
|
||||
filtered_torrents.append(torrent)
|
||||
media_type_enum = None
|
||||
if media_type:
|
||||
media_type_enum = MediaType.from_agent(media_type)
|
||||
if not media_type_enum:
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv'"
|
||||
|
||||
filtered_torrents = await search_chain.async_search_by_id(
|
||||
tmdbid=tmdb_id,
|
||||
doubanid=douban_id,
|
||||
mtype=media_type_enum,
|
||||
area=area or "title",
|
||||
sites=sites,
|
||||
cache_local=False,
|
||||
)
|
||||
|
||||
# 获取站点信息
|
||||
all_indexers = await SitesHelper().async_get_indexers()
|
||||
all_sites = [{"id": indexer.get("id"), "name": indexer.get("name")} for indexer in (all_indexers or [])]
|
||||
|
||||
if sites:
|
||||
search_site_ids = sites
|
||||
else:
|
||||
configured_sites = SystemConfigOper().get(SystemConfigKey.IndexerSites)
|
||||
search_site_ids = configured_sites if configured_sites else []
|
||||
|
||||
if filtered_torrents:
|
||||
# 限制最多50条结果
|
||||
total_count = len(filtered_torrents)
|
||||
limited_torrents = filtered_torrents[:50]
|
||||
# 精简字段,只保留关键信息
|
||||
simplified_torrents = []
|
||||
for t in limited_torrents:
|
||||
simplified = {}
|
||||
# 精简 torrent_info
|
||||
if t.torrent_info:
|
||||
simplified["torrent_info"] = {
|
||||
"title": t.torrent_info.title,
|
||||
"size": t.torrent_info.size,
|
||||
"seeders": t.torrent_info.seeders,
|
||||
"peers": t.torrent_info.peers,
|
||||
"site_name": t.torrent_info.site_name,
|
||||
"enclosure": t.torrent_info.enclosure,
|
||||
"page_url": t.torrent_info.page_url,
|
||||
"volume_factor": t.torrent_info.volume_factor,
|
||||
"pubdate": t.torrent_info.pubdate
|
||||
}
|
||||
# 精简 media_info
|
||||
if t.media_info:
|
||||
simplified["media_info"] = {
|
||||
"title": t.media_info.title,
|
||||
"en_title": t.media_info.en_title,
|
||||
"year": t.media_info.year,
|
||||
"type": t.media_info.type.value if t.media_info.type else None,
|
||||
"season": t.media_info.season,
|
||||
"tmdb_id": t.media_info.tmdb_id
|
||||
}
|
||||
# 精简 meta_info
|
||||
if t.meta_info:
|
||||
simplified["meta_info"] = {
|
||||
"name": t.meta_info.name,
|
||||
"cn_name": t.meta_info.cn_name,
|
||||
"en_name": t.meta_info.en_name,
|
||||
"year": t.meta_info.year,
|
||||
"type": t.meta_info.type.value if t.meta_info.type else None,
|
||||
"begin_season": t.meta_info.begin_season
|
||||
}
|
||||
simplified_torrents.append(simplified)
|
||||
result_json = json.dumps(simplified_torrents, ensure_ascii=False, indent=2)
|
||||
# 如果结果被裁剪,添加提示信息
|
||||
if total_count > 50:
|
||||
return f"注意:搜索结果共找到 {total_count} 条,为节省上下文空间,仅显示前 50 条结果。\n\n{result_json}"
|
||||
await search_chain.async_save_cache(filtered_torrents, SEARCH_RESULT_CACHE_FILE)
|
||||
result_json = json.dumps({
|
||||
"total_count": len(filtered_torrents),
|
||||
"message": "搜索完成。请使用 get_search_results 工具获取搜索结果。",
|
||||
"all_sites": all_sites,
|
||||
"search_site_ids": search_site_ids,
|
||||
"filter_options": build_filter_options(filtered_torrents),
|
||||
}, ensure_ascii=False, indent=2)
|
||||
return result_json
|
||||
else:
|
||||
return f"未找到相关种子资源: {title}"
|
||||
media_id = f"TMDB={tmdb_id}" if tmdb_id else f"豆瓣={douban_id}"
|
||||
result_json = json.dumps({
|
||||
"message": f"未找到相关种子资源: {media_id}",
|
||||
"all_sites": all_sites,
|
||||
"search_site_ids": search_site_ids,
|
||||
}, ensure_ascii=False, indent=2)
|
||||
return result_json
|
||||
except Exception as e:
|
||||
error_message = f"搜索种子时发生错误: {str(e)}"
|
||||
logger.error(f"搜索种子失败: {e}", exc_info=True)
|
||||
|
||||
@@ -1,22 +1,26 @@
|
||||
"""搜索网络内容工具"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from typing import Optional, Type
|
||||
from typing import Optional, Type, List, Dict
|
||||
|
||||
import httpx
|
||||
from ddgs import DDGS
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.core.config import settings
|
||||
from app.log import logger
|
||||
from app.utils.http import AsyncRequestUtils
|
||||
|
||||
# 搜索超时时间(秒)
|
||||
SEARCH_TIMEOUT = 20
|
||||
|
||||
|
||||
class SearchWebInput(BaseModel):
|
||||
"""搜索网络内容工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
query: str = Field(..., description="The search query string to search for on the web")
|
||||
max_results: Optional[int] = Field(5, description="Maximum number of search results to return (default: 5, max: 10)")
|
||||
max_results: Optional[int] = Field(5,
|
||||
description="Maximum number of search results to return (default: 5, max: 10)")
|
||||
|
||||
|
||||
class SearchWebTool(MoviePilotTool):
|
||||
@@ -33,151 +37,137 @@ class SearchWebTool(MoviePilotTool):
|
||||
async def run(self, query: str, max_results: Optional[int] = 5, **kwargs) -> str:
|
||||
"""
|
||||
执行网络搜索
|
||||
|
||||
Args:
|
||||
query: 搜索查询字符串
|
||||
max_results: 最大返回结果数(默认5,最大10)
|
||||
|
||||
Returns:
|
||||
格式化的搜索结果JSON字符串
|
||||
"""
|
||||
logger.info(f"执行工具: {self.name}, 参数: query={query}, max_results={max_results}")
|
||||
|
||||
try:
|
||||
# 限制最大结果数
|
||||
max_results = min(max(1, max_results or 5), 10)
|
||||
|
||||
# 使用DuckDuckGo API进行搜索
|
||||
search_results = await self._search_duckduckgo_api(query, max_results)
|
||||
|
||||
if not search_results:
|
||||
results = []
|
||||
|
||||
# 1. 优先使用 Tavily (如果配置了 API Key)
|
||||
if settings.TAVILY_API_KEY:
|
||||
logger.info("使用 Tavily 进行搜索...")
|
||||
results = await self._search_tavily(query, max_results)
|
||||
|
||||
# 2. 如果没有结果或未配置 Tavily,使用 DuckDuckGo
|
||||
if not results:
|
||||
logger.info("使用 DuckDuckGo 进行搜索...")
|
||||
results = await self._search_duckduckgo(query, max_results)
|
||||
|
||||
if not results:
|
||||
return f"未找到与 '{query}' 相关的搜索结果"
|
||||
|
||||
# 裁剪结果以避免占用过多上下文
|
||||
formatted_results = self._format_and_truncate_results(search_results, max_results)
|
||||
|
||||
result_json = json.dumps(formatted_results, ensure_ascii=False, indent=2)
|
||||
return result_json
|
||||
|
||||
|
||||
# 格式化并裁剪结果
|
||||
formatted_results = self._format_and_truncate_results(results, max_results)
|
||||
return json.dumps(formatted_results, ensure_ascii=False, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
error_message = f"搜索网络内容失败: {str(e)}"
|
||||
logger.error(f"搜索网络内容失败: {e}", exc_info=True)
|
||||
return error_message
|
||||
|
||||
@staticmethod
|
||||
async def _search_duckduckgo_api(query: str, max_results: int) -> list:
|
||||
"""
|
||||
使用DuckDuckGo API进行搜索
|
||||
|
||||
Args:
|
||||
query: 搜索查询
|
||||
max_results: 最大结果数
|
||||
|
||||
Returns:
|
||||
搜索结果列表
|
||||
"""
|
||||
async def _search_tavily(query: str, max_results: int) -> List[Dict]:
|
||||
"""使用 Tavily API 进行搜索"""
|
||||
try:
|
||||
# DuckDuckGo Instant Answer API
|
||||
api_url = "https://api.duckduckgo.com/"
|
||||
params = {
|
||||
"q": query,
|
||||
"format": "json",
|
||||
"no_html": "1",
|
||||
"skip_disambig": "1"
|
||||
}
|
||||
|
||||
# 使用代理(如果配置了)
|
||||
http_utils = AsyncRequestUtils(
|
||||
proxies=settings.PROXY,
|
||||
timeout=10
|
||||
)
|
||||
|
||||
data = await http_utils.get_json(api_url, params=params)
|
||||
|
||||
results = []
|
||||
|
||||
if data:
|
||||
# 处理AbstractText(摘要)
|
||||
if data.get("AbstractText"):
|
||||
async with httpx.AsyncClient(timeout=SEARCH_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
"https://api.tavily.com/search",
|
||||
json={
|
||||
"api_key": settings.TAVILY_API_KEY,
|
||||
"query": query,
|
||||
"search_depth": "basic",
|
||||
"max_results": max_results,
|
||||
"include_answer": False,
|
||||
"include_images": False,
|
||||
"include_raw_content": False,
|
||||
}
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
results = []
|
||||
for result in data.get("results", []):
|
||||
results.append({
|
||||
"title": data.get("Heading", query),
|
||||
"snippet": data.get("AbstractText", ""),
|
||||
"url": data.get("AbstractURL", ""),
|
||||
"source": "DuckDuckGo Abstract"
|
||||
'title': result.get('title', ''),
|
||||
'snippet': result.get('content', ''),
|
||||
'url': result.get('url', ''),
|
||||
'source': 'Tavily'
|
||||
})
|
||||
|
||||
# 处理RelatedTopics(相关主题)
|
||||
related_topics = data.get("RelatedTopics", [])
|
||||
for topic in related_topics[:max_results - len(results)]:
|
||||
if isinstance(topic, dict):
|
||||
text = topic.get("Text", "")
|
||||
first_url = topic.get("FirstURL", "")
|
||||
if text and first_url:
|
||||
# 提取标题(通常在" - "之前)
|
||||
title = text.split(" - ")[0] if " - " in text else text[:100]
|
||||
snippet = text
|
||||
|
||||
results.append({
|
||||
"title": title.strip(),
|
||||
"snippet": snippet,
|
||||
"url": first_url,
|
||||
"source": "DuckDuckGo Related"
|
||||
})
|
||||
|
||||
# 处理Results(搜索结果)
|
||||
api_results = data.get("Results", [])
|
||||
for result in api_results[:max_results - len(results)]:
|
||||
if isinstance(result, dict):
|
||||
title = result.get("Text", "")
|
||||
url = result.get("FirstURL", "")
|
||||
if title and url:
|
||||
results.append({
|
||||
"title": title,
|
||||
"snippet": result.get("Text", ""),
|
||||
"url": url,
|
||||
"source": "DuckDuckGo Results"
|
||||
})
|
||||
|
||||
return results[:max_results]
|
||||
|
||||
return results
|
||||
except Exception as e:
|
||||
logger.warning(f"DuckDuckGo API搜索失败: {e}")
|
||||
logger.warning(f"Tavily 搜索失败: {e}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _format_and_truncate_results(results: list, max_results: int) -> dict:
|
||||
"""
|
||||
格式化并裁剪搜索结果以避免占用过多上下文
|
||||
|
||||
Args:
|
||||
results: 原始搜索结果列表
|
||||
max_results: 最大结果数
|
||||
|
||||
Returns:
|
||||
格式化后的结果字典
|
||||
"""
|
||||
def _get_proxy_url(proxy_setting) -> Optional[str]:
|
||||
"""从代理设置中提取代理URL"""
|
||||
if not proxy_setting:
|
||||
return None
|
||||
if isinstance(proxy_setting, dict):
|
||||
return proxy_setting.get('http') or proxy_setting.get('https')
|
||||
return proxy_setting
|
||||
|
||||
async def _search_duckduckgo(self, query: str, max_results: int) -> List[Dict]:
|
||||
"""使用 duckduckgo-search (DDGS) 进行搜索"""
|
||||
try:
|
||||
def sync_search():
|
||||
results = []
|
||||
ddgs_kwargs = {
|
||||
'timeout': SEARCH_TIMEOUT
|
||||
}
|
||||
proxy_url = self._get_proxy_url(settings.PROXY)
|
||||
if proxy_url:
|
||||
ddgs_kwargs['proxy'] = proxy_url
|
||||
|
||||
try:
|
||||
with DDGS(**ddgs_kwargs) as ddgs:
|
||||
ddgs_gen = ddgs.text(
|
||||
query,
|
||||
max_results=max_results
|
||||
)
|
||||
if ddgs_gen:
|
||||
for result in ddgs_gen:
|
||||
results.append({
|
||||
'title': result.get('title', ''),
|
||||
'snippet': result.get('body', ''),
|
||||
'url': result.get('href', ''),
|
||||
'source': 'DuckDuckGo'
|
||||
})
|
||||
except Exception as err:
|
||||
logger.warning(f"DuckDuckGo search process failed: {err}")
|
||||
return results
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(None, sync_search)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"DuckDuckGo 搜索失败: {e}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _format_and_truncate_results(results: List[Dict], max_results: int) -> Dict:
|
||||
"""格式化并裁剪搜索结果"""
|
||||
formatted = {
|
||||
"total_results": len(results),
|
||||
"results": []
|
||||
}
|
||||
|
||||
# 限制结果数量
|
||||
limited_results = results[:max_results]
|
||||
|
||||
for idx, result in enumerate(limited_results, 1):
|
||||
title = result.get("title", "")[:200] # 限制标题长度
|
||||
|
||||
for idx, result in enumerate(results[:max_results], 1):
|
||||
title = result.get("title", "")[:200]
|
||||
snippet = result.get("snippet", "")
|
||||
url = result.get("url", "")
|
||||
source = result.get("source", "Unknown")
|
||||
|
||||
# 裁剪摘要,避免过长
|
||||
max_snippet_length = 300 # 每个摘要最多300字符
|
||||
|
||||
# 裁剪摘要
|
||||
max_snippet_length = 500 # 增加到500字符,提供更多上下文
|
||||
if len(snippet) > max_snippet_length:
|
||||
snippet = snippet[:max_snippet_length] + "..."
|
||||
|
||||
# 清理文本,移除多余的空白字符
|
||||
|
||||
# 清理文本
|
||||
snippet = re.sub(r'\s+', ' ', snippet).strip()
|
||||
|
||||
|
||||
formatted["results"].append({
|
||||
"rank": idx,
|
||||
"title": title,
|
||||
@@ -185,9 +175,8 @@ class SearchWebTool(MoviePilotTool):
|
||||
"url": url,
|
||||
"source": source
|
||||
})
|
||||
|
||||
# 添加提示信息
|
||||
|
||||
if len(results) > max_results:
|
||||
formatted["note"] = f"注意:共找到 {len(results)} 条结果,为节省上下文空间,仅显示前 {max_results} 条结果。"
|
||||
|
||||
formatted["note"] = f"仅显示前 {max_results} 条结果。"
|
||||
|
||||
return formatted
|
||||
|
||||
@@ -8,53 +8,31 @@ from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.site import SiteChain
|
||||
from app.db.site_oper import SiteOper
|
||||
from app.log import logger
|
||||
from app.utils.string import StringUtils
|
||||
|
||||
|
||||
class TestSiteInput(BaseModel):
|
||||
"""测试站点连通性工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
site_identifier: str = Field(..., description="Site identifier: can be site ID (integer as string), site name, or site domain/URL")
|
||||
site_identifier: int = Field(..., description="Site ID to test (can be obtained from query_sites tool)")
|
||||
|
||||
|
||||
class TestSiteTool(MoviePilotTool):
|
||||
name: str = "test_site"
|
||||
description: str = "Test site connectivity and availability. This will check if a site is accessible and can be logged in. Accepts site ID, site name, or site domain/URL as identifier."
|
||||
description: str = "Test site connectivity and availability. This will check if a site is accessible and can be logged in. Accepts site ID only."
|
||||
args_schema: Type[BaseModel] = TestSiteInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据测试参数生成友好的提示消息"""
|
||||
site_identifier = kwargs.get("site_identifier", "")
|
||||
site_identifier = kwargs.get("site_identifier")
|
||||
return f"正在测试站点连通性: {site_identifier}"
|
||||
|
||||
async def run(self, site_identifier: str, **kwargs) -> str:
|
||||
async def run(self, site_identifier: int, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: site_identifier={site_identifier}")
|
||||
|
||||
try:
|
||||
site_oper = SiteOper()
|
||||
site_chain = SiteChain()
|
||||
|
||||
# 尝试解析为站点ID
|
||||
site = None
|
||||
if site_identifier.isdigit():
|
||||
# 如果是数字,尝试作为站点ID查询
|
||||
site = await site_oper.async_get(int(site_identifier))
|
||||
|
||||
# 如果不是ID或ID查询失败,尝试按名称或域名查询
|
||||
if not site:
|
||||
# 尝试按名称查询
|
||||
sites = await site_oper.async_list()
|
||||
for s in sites:
|
||||
if (site_identifier.lower() in (s.name or "").lower()) or \
|
||||
(site_identifier.lower() in (s.domain or "").lower()):
|
||||
site = s
|
||||
break
|
||||
|
||||
# 如果还是没找到,尝试从URL提取域名
|
||||
if not site:
|
||||
domain = StringUtils.get_url_domain(site_identifier)
|
||||
if domain:
|
||||
site = await site_oper.async_get_by_domain(domain)
|
||||
site = await site_oper.async_get(site_identifier)
|
||||
|
||||
if not site:
|
||||
return f"未找到站点:{site_identifier},请使用 query_sites 工具查询可用的站点"
|
||||
|
||||
@@ -18,7 +18,7 @@ class TransferFileInput(BaseModel):
|
||||
storage: Optional[str] = Field("local", description="Storage type of the source file (default: 'local', can be 'smb', 'alist', etc.)")
|
||||
target_path: Optional[str] = Field(None, description="Target path for the transferred file/directory (optional, uses default library path if not specified)")
|
||||
target_storage: Optional[str] = Field(None, description="Target storage type (optional, uses default storage if not specified)")
|
||||
media_type: Optional[str] = Field(None, description="Media type: '电影' for films, '电视剧' for television series (optional, will be auto-detected if not specified)")
|
||||
media_type: Optional[str] = Field(None, description="Allowed values: movie, tv")
|
||||
tmdbid: Optional[int] = Field(None, description="TMDB ID for precise media identification (optional but recommended for accuracy)")
|
||||
doubanid: Optional[str] = Field(None, description="Douban ID for media identification (optional)")
|
||||
season: Optional[int] = Field(None, description="Season number for TV shows (optional)")
|
||||
@@ -91,11 +91,10 @@ class TransferFileTool(MoviePilotTool):
|
||||
target_path_obj = Path(target_path)
|
||||
|
||||
# 处理媒体类型
|
||||
mtype = None
|
||||
media_type_enum = None
|
||||
if media_type:
|
||||
try:
|
||||
mtype = MediaType(media_type)
|
||||
except ValueError:
|
||||
media_type_enum = MediaType.from_agent(media_type)
|
||||
if not media_type_enum:
|
||||
return f"错误:无效的媒体类型 '{media_type}',支持的类型:'movie', 'tv'"
|
||||
|
||||
# 调用整理方法
|
||||
@@ -106,7 +105,7 @@ class TransferFileTool(MoviePilotTool):
|
||||
target_path=target_path_obj,
|
||||
tmdbid=tmdbid,
|
||||
doubanid=doubanid,
|
||||
mtype=mtype,
|
||||
mtype=media_type_enum,
|
||||
season=season,
|
||||
transfer_type=transfer_type,
|
||||
background=background
|
||||
|
||||
@@ -17,7 +17,7 @@ from app.utils.string import StringUtils
|
||||
class UpdateSiteInput(BaseModel):
|
||||
"""更新站点工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
site_id: int = Field(..., description="The ID of the site to update")
|
||||
site_id: int = Field(..., description="The ID of the site to update (can be obtained from query_sites tool)")
|
||||
name: Optional[str] = Field(None, description="Site name (optional)")
|
||||
url: Optional[str] = Field(None, description="Site URL (optional, will be automatically formatted)")
|
||||
pri: Optional[int] = Field(None, description="Site priority (optional, smaller value = higher priority, e.g., pri=1 has higher priority than pri=10)")
|
||||
|
||||
@@ -8,13 +8,12 @@ from app.agent.tools.base import MoviePilotTool
|
||||
from app.chain.site import SiteChain
|
||||
from app.db.site_oper import SiteOper
|
||||
from app.log import logger
|
||||
from app.utils.string import StringUtils
|
||||
|
||||
|
||||
class UpdateSiteCookieInput(BaseModel):
|
||||
"""更新站点Cookie和UA工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
site_identifier: str = Field(..., description="Site identifier: can be site ID (integer as string), site name, or site domain/URL")
|
||||
site_identifier: int = Field(..., description="Site ID to update Cookie and User-Agent for (can be obtained from query_sites tool)")
|
||||
username: str = Field(..., description="Site login username")
|
||||
password: str = Field(..., description="Site login password")
|
||||
two_step_code: Optional[str] = Field(None, description="Two-step verification code or secret key (optional, required for sites with 2FA enabled)")
|
||||
@@ -22,12 +21,12 @@ class UpdateSiteCookieInput(BaseModel):
|
||||
|
||||
class UpdateSiteCookieTool(MoviePilotTool):
|
||||
name: str = "update_site_cookie"
|
||||
description: str = "Update site Cookie and User-Agent by logging in with username and password. This tool can automatically obtain and update the site's authentication credentials. Supports two-step verification for sites that require it. Accepts site ID, site name, or site domain/URL as identifier."
|
||||
description: str = "Update site Cookie and User-Agent by logging in with username and password. This tool can automatically obtain and update the site's authentication credentials. Supports two-step verification for sites that require it. Accepts site ID only."
|
||||
args_schema: Type[BaseModel] = UpdateSiteCookieInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据更新参数生成友好的提示消息"""
|
||||
site_identifier = kwargs.get("site_identifier", "")
|
||||
site_identifier = kwargs.get("site_identifier")
|
||||
username = kwargs.get("username", "")
|
||||
two_step_code = kwargs.get("two_step_code")
|
||||
|
||||
@@ -37,35 +36,14 @@ class UpdateSiteCookieTool(MoviePilotTool):
|
||||
|
||||
return message
|
||||
|
||||
async def run(self, site_identifier: str, username: str, password: str,
|
||||
async def run(self, site_identifier: int, username: str, password: str,
|
||||
two_step_code: Optional[str] = None, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: site_identifier={site_identifier}, username={username}")
|
||||
|
||||
try:
|
||||
site_oper = SiteOper()
|
||||
site_chain = SiteChain()
|
||||
|
||||
# 尝试解析为站点ID
|
||||
site = None
|
||||
if site_identifier.isdigit():
|
||||
# 如果是数字,尝试作为站点ID查询
|
||||
site = await site_oper.async_get(int(site_identifier))
|
||||
|
||||
# 如果不是ID或ID查询失败,尝试按名称或域名查询
|
||||
if not site:
|
||||
# 尝试按名称查询
|
||||
sites = await site_oper.async_list()
|
||||
for s in sites:
|
||||
if (site_identifier.lower() in (s.name or "").lower()) or \
|
||||
(site_identifier.lower() in (s.domain or "").lower()):
|
||||
site = s
|
||||
break
|
||||
|
||||
# 如果还是没找到,尝试从URL提取域名
|
||||
if not site:
|
||||
domain = StringUtils.get_url_domain(site_identifier)
|
||||
if domain:
|
||||
site = await site_oper.async_get_by_domain(domain)
|
||||
site = await site_oper.async_get(site_identifier)
|
||||
|
||||
if not site:
|
||||
return f"未找到站点:{site_identifier},请使用 query_sites 工具查询可用的站点"
|
||||
|
||||
@@ -16,7 +16,7 @@ from app.schemas.types import EventType
|
||||
class UpdateSubscribeInput(BaseModel):
|
||||
"""更新订阅工具的输入参数模型"""
|
||||
explanation: str = Field(..., description="Clear explanation of why this tool is being used in the current context")
|
||||
subscribe_id: int = Field(..., description="The ID of the subscription to update")
|
||||
subscribe_id: int = Field(..., description="The ID of the subscription to update (can be obtained from query_subscribes tool)")
|
||||
name: Optional[str] = Field(None, description="Subscription name/title (optional)")
|
||||
year: Optional[str] = Field(None, description="Release year (optional)")
|
||||
season: Optional[int] = Field(None, description="Season number for TV shows (optional)")
|
||||
@@ -29,7 +29,7 @@ class UpdateSubscribeInput(BaseModel):
|
||||
include: Optional[str] = Field(None, description="Include filter as regular expression (optional)")
|
||||
exclude: Optional[str] = Field(None, description="Exclude filter as regular expression (optional)")
|
||||
filter: Optional[str] = Field(None, description="Filter rule as regular expression (optional)")
|
||||
state: Optional[str] = Field(None, description="Subscription state: 'R' for enabled, 'P' for disabled, 'S' for paused (optional)")
|
||||
state: Optional[str] = Field(None, description="Subscription state: 'R' for enabled, 'P' for pending, 'S' for paused (optional)")
|
||||
sites: Optional[List[int]] = Field(None, description="List of site IDs to search from (optional)")
|
||||
downloader: Optional[str] = Field(None, description="Downloader name (optional)")
|
||||
save_path: Optional[str] = Field(None, description="Save path for downloaded files (optional)")
|
||||
|
||||
52
app/agent/tools/impl/write_file.py
Normal file
52
app/agent/tools/impl/write_file.py
Normal file
@@ -0,0 +1,52 @@
|
||||
"""文件写入工具"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Optional, Type
|
||||
|
||||
from anyio import Path as AsyncPath
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.agent.tools.base import MoviePilotTool
|
||||
from app.log import logger
|
||||
|
||||
|
||||
class WriteFileInput(BaseModel):
|
||||
"""Input parameters for write file tool"""
|
||||
file_path: str = Field(..., description="The absolute path of the file to write")
|
||||
content: str = Field(..., description="The content to write into the file")
|
||||
|
||||
|
||||
class WriteFileTool(MoviePilotTool):
|
||||
name: str = "write_file"
|
||||
description: str = "Write full content to a file. If the file already exists, it will be overwritten. Automatically creates parent directories if they don't exist."
|
||||
args_schema: Type[BaseModel] = WriteFileInput
|
||||
|
||||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||||
"""根据参数生成友好的提示消息"""
|
||||
file_path = kwargs.get("file_path", "")
|
||||
file_name = Path(file_path).name if file_path else "未知文件"
|
||||
return f"正在写入文件: {file_name}"
|
||||
|
||||
async def run(self, file_path: str, content: str, **kwargs) -> str:
|
||||
logger.info(f"执行工具: {self.name}, 参数: file_path={file_path}")
|
||||
|
||||
try:
|
||||
path = AsyncPath(file_path)
|
||||
|
||||
if await path.exists() and not await path.is_file():
|
||||
return f"错误:{file_path} 路径已存在但不是一个文件"
|
||||
|
||||
# 自动创建父目录
|
||||
await path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 写入文件
|
||||
await path.write_text(content, encoding="utf-8")
|
||||
|
||||
logger.info(f"成功写入文件 {file_path}")
|
||||
return f"成功写入文件 {file_path}"
|
||||
|
||||
except PermissionError:
|
||||
return f"错误:没有权限写入 {file_path}"
|
||||
except Exception as e:
|
||||
logger.error(f"写入文件 {file_path} 时发生错误: {str(e)}", exc_info=True)
|
||||
return f"操作失败: {str(e)}"
|
||||
@@ -1,18 +1,15 @@
|
||||
"""MoviePilot工具管理器
|
||||
用于HTTP API调用工具
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from app.agent import ConversationMemoryManager
|
||||
from app.agent.tools.factory import MoviePilotToolFactory
|
||||
from app.log import logger
|
||||
|
||||
|
||||
class ToolDefinition:
|
||||
"""工具定义"""
|
||||
"""
|
||||
工具定义
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, description: str, input_schema: Dict[str, Any]):
|
||||
self.name = name
|
||||
@@ -21,12 +18,14 @@ class ToolDefinition:
|
||||
|
||||
|
||||
class MoviePilotToolsManager:
|
||||
"""MoviePilot工具管理器(用于HTTP API)"""
|
||||
"""
|
||||
MoviePilot工具管理器(用于HTTP API)
|
||||
"""
|
||||
|
||||
def __init__(self, user_id: str = "api_user", session_id: str = uuid.uuid4()):
|
||||
"""
|
||||
初始化工具管理器
|
||||
|
||||
|
||||
Args:
|
||||
user_id: 用户ID
|
||||
session_id: 会话ID
|
||||
@@ -34,11 +33,12 @@ class MoviePilotToolsManager:
|
||||
self.user_id = user_id
|
||||
self.session_id = session_id
|
||||
self.tools: List[Any] = []
|
||||
self.memory_manager = ConversationMemoryManager()
|
||||
self._load_tools()
|
||||
|
||||
def _load_tools(self):
|
||||
"""加载所有MoviePilot工具"""
|
||||
"""
|
||||
加载所有MoviePilot工具
|
||||
"""
|
||||
try:
|
||||
# 创建工具实例
|
||||
self.tools = MoviePilotToolFactory.create_tools(
|
||||
@@ -47,8 +47,7 @@ class MoviePilotToolsManager:
|
||||
channel=None,
|
||||
source="api",
|
||||
username="API Client",
|
||||
callback_handler=None,
|
||||
memory_mananger=None,
|
||||
stream_handler=None,
|
||||
)
|
||||
logger.info(f"成功加载 {len(self.tools)} 个工具")
|
||||
except Exception as e:
|
||||
@@ -58,40 +57,38 @@ class MoviePilotToolsManager:
|
||||
def list_tools(self) -> List[ToolDefinition]:
|
||||
"""
|
||||
列出所有可用的工具
|
||||
|
||||
|
||||
Returns:
|
||||
工具定义列表
|
||||
"""
|
||||
tools_list = []
|
||||
for tool in self.tools:
|
||||
# 获取工具的输入参数模型
|
||||
args_schema = getattr(tool, 'args_schema', None)
|
||||
args_schema = getattr(tool, "args_schema", None)
|
||||
if args_schema:
|
||||
# 将Pydantic模型转换为JSON Schema
|
||||
input_schema = self._convert_to_json_schema(args_schema)
|
||||
else:
|
||||
# 如果没有args_schema,使用基本信息
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
input_schema = {"type": "object", "properties": {}, "required": []}
|
||||
|
||||
tools_list.append(ToolDefinition(
|
||||
name=tool.name,
|
||||
description=tool.description or "",
|
||||
input_schema=input_schema
|
||||
))
|
||||
tools_list.append(
|
||||
ToolDefinition(
|
||||
name=tool.name,
|
||||
description=tool.description or "",
|
||||
input_schema=input_schema,
|
||||
)
|
||||
)
|
||||
|
||||
return tools_list
|
||||
|
||||
def get_tool(self, tool_name: str) -> Optional[Any]:
|
||||
"""
|
||||
获取指定工具实例
|
||||
|
||||
|
||||
Args:
|
||||
tool_name: 工具名称
|
||||
|
||||
|
||||
Returns:
|
||||
工具实例,如果未找到返回None
|
||||
"""
|
||||
@@ -101,22 +98,88 @@ class MoviePilotToolsManager:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _normalize_arguments(tool_instance: Any, arguments: Dict[str, Any]) -> Dict[str, Any]:
|
||||
def _resolve_field_schema(field_info: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
解析字段schema,兼容 Optional[T] 生成的 anyOf 结构
|
||||
"""
|
||||
if field_info.get("type"):
|
||||
return field_info
|
||||
|
||||
any_of = field_info.get("anyOf")
|
||||
if not any_of:
|
||||
return field_info
|
||||
|
||||
for type_option in any_of:
|
||||
if type_option.get("type") and type_option["type"] != "null":
|
||||
merged = dict(type_option)
|
||||
if "description" not in merged and field_info.get("description"):
|
||||
merged["description"] = field_info["description"]
|
||||
if "default" not in merged and "default" in field_info:
|
||||
merged["default"] = field_info["default"]
|
||||
return merged
|
||||
|
||||
return field_info
|
||||
|
||||
@staticmethod
|
||||
def _normalize_scalar_value(field_type: Optional[str], value: Any, key: str) -> Any:
|
||||
"""
|
||||
根据字段类型规范化单个值
|
||||
"""
|
||||
if field_type == "integer" and isinstance(value, str):
|
||||
try:
|
||||
return int(value)
|
||||
except (ValueError, TypeError):
|
||||
logger.warning(f"无法将参数 {key}='{value}' 转换为整数,返回 None")
|
||||
return None
|
||||
if field_type == "number" and isinstance(value, str):
|
||||
try:
|
||||
return float(value)
|
||||
except (ValueError, TypeError):
|
||||
logger.warning(f"无法将参数 {key}='{value}' 转换为浮点数,返回 None")
|
||||
return None
|
||||
if field_type == "boolean":
|
||||
if isinstance(value, str):
|
||||
return value.lower() in ("true", "1", "yes", "on")
|
||||
if isinstance(value, (int, float)):
|
||||
return value != 0
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
return True
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _parse_array_string(value: str, key: str, item_type: str = "string") -> list:
|
||||
"""
|
||||
将逗号分隔的字符串解析为列表,并根据 item_type 转换元素类型
|
||||
"""
|
||||
trimmed = value.strip()
|
||||
if not trimmed:
|
||||
return []
|
||||
return [
|
||||
MoviePilotToolsManager._normalize_scalar_value(item_type, item.strip(), key)
|
||||
for item in trimmed.split(",")
|
||||
if item.strip()
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _normalize_arguments(
|
||||
tool_instance: Any, arguments: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
根据工具的参数schema规范化参数类型
|
||||
|
||||
|
||||
Args:
|
||||
tool_instance: 工具实例
|
||||
arguments: 原始参数
|
||||
|
||||
|
||||
Returns:
|
||||
规范化后的参数
|
||||
"""
|
||||
# 获取工具的参数schema
|
||||
args_schema = getattr(tool_instance, 'args_schema', None)
|
||||
args_schema = getattr(tool_instance, "args_schema", None)
|
||||
if not args_schema:
|
||||
return arguments
|
||||
|
||||
|
||||
# 获取schema中的字段定义
|
||||
try:
|
||||
schema = args_schema.model_json_schema()
|
||||
@@ -124,7 +187,7 @@ class MoviePilotToolsManager:
|
||||
except Exception as e:
|
||||
logger.warning(f"获取工具schema失败: {e}")
|
||||
return arguments
|
||||
|
||||
|
||||
# 规范化参数
|
||||
normalized = {}
|
||||
for key, value in arguments.items():
|
||||
@@ -132,91 +195,80 @@ class MoviePilotToolsManager:
|
||||
# 参数不在schema中,保持原样
|
||||
normalized[key] = value
|
||||
continue
|
||||
|
||||
field_info = properties[key]
|
||||
|
||||
field_info = MoviePilotToolsManager._resolve_field_schema(properties[key])
|
||||
field_type = field_info.get("type")
|
||||
|
||||
# 处理 anyOf 类型(例如 Optional[int] 会生成 anyOf)
|
||||
any_of = field_info.get("anyOf")
|
||||
if any_of and not field_type:
|
||||
# 从 anyOf 中提取实际类型
|
||||
for type_option in any_of:
|
||||
if "type" in type_option and type_option["type"] != "null":
|
||||
field_type = type_option["type"]
|
||||
break
|
||||
|
||||
|
||||
# 数组类型:将字符串解析为列表
|
||||
if field_type == "array" and isinstance(value, str):
|
||||
item_type = field_info.get("items", {}).get("type", "string")
|
||||
normalized[key] = MoviePilotToolsManager._parse_array_string(
|
||||
value, key, item_type
|
||||
)
|
||||
continue
|
||||
|
||||
# 根据类型进行转换
|
||||
if field_type == "integer" and isinstance(value, str):
|
||||
try:
|
||||
normalized[key] = int(value)
|
||||
except (ValueError, TypeError):
|
||||
logger.warning(f"无法将参数 {key}='{value}' 转换为整数,保持原值")
|
||||
normalized[key] = value
|
||||
elif field_type == "number" and isinstance(value, str):
|
||||
try:
|
||||
normalized[key] = float(value)
|
||||
except (ValueError, TypeError):
|
||||
logger.warning(f"无法将参数 {key}='{value}' 转换为浮点数,保持原值")
|
||||
normalized[key] = value
|
||||
elif field_type == "boolean" and isinstance(value, str):
|
||||
# 转换字符串为布尔值
|
||||
normalized[key] = value.lower() in ("true", "1", "yes", "on")
|
||||
else:
|
||||
# 其他类型保持原样
|
||||
normalized[key] = value
|
||||
|
||||
normalized[key] = MoviePilotToolsManager._normalize_scalar_value(
|
||||
field_type, value, key
|
||||
)
|
||||
|
||||
return normalized
|
||||
|
||||
async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str:
|
||||
"""
|
||||
调用工具
|
||||
|
||||
|
||||
Args:
|
||||
tool_name: 工具名称
|
||||
arguments: 工具参数
|
||||
|
||||
|
||||
Returns:
|
||||
工具执行结果(字符串)
|
||||
"""
|
||||
tool_instance = self.get_tool(tool_name)
|
||||
|
||||
if not tool_instance:
|
||||
error_msg = json.dumps({
|
||||
"error": f"工具 '{tool_name}' 未找到"
|
||||
}, ensure_ascii=False)
|
||||
error_msg = json.dumps(
|
||||
{"error": f"工具 '{tool_name}' 未找到"}, ensure_ascii=False
|
||||
)
|
||||
return error_msg
|
||||
|
||||
try:
|
||||
# 规范化参数类型
|
||||
normalized_arguments = self._normalize_arguments(tool_instance, arguments)
|
||||
|
||||
|
||||
# 调用工具的run方法
|
||||
result = await tool_instance.run(**normalized_arguments)
|
||||
|
||||
# 确保返回字符串
|
||||
if isinstance(result, str):
|
||||
formated_result = result
|
||||
elif isinstance(result, int, float):
|
||||
elif isinstance(result, (int, float)):
|
||||
formated_result = str(result)
|
||||
else:
|
||||
formated_result = json.dumps(result, ensure_ascii=False, indent=2)
|
||||
try:
|
||||
formated_result = json.dumps(result, ensure_ascii=False, indent=2)
|
||||
except Exception as e:
|
||||
logger.warning(f"结果转换为JSON失败: {e}, 使用字符串表示")
|
||||
formated_result = str(result)
|
||||
|
||||
return formated_result
|
||||
except Exception as e:
|
||||
logger.error(f"调用工具 {tool_name} 时发生错误: {e}", exc_info=True)
|
||||
error_msg = json.dumps({
|
||||
"error": f"调用工具 '{tool_name}' 时发生错误: {str(e)}"
|
||||
}, ensure_ascii=False)
|
||||
error_msg = json.dumps(
|
||||
{"error": f"调用工具 '{tool_name}' 时发生错误: {str(e)}"},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
return error_msg
|
||||
|
||||
@staticmethod
|
||||
def _convert_to_json_schema(args_schema: Any) -> Dict[str, Any]:
|
||||
"""
|
||||
将Pydantic模型转换为JSON Schema
|
||||
|
||||
|
||||
Args:
|
||||
args_schema: Pydantic模型类
|
||||
|
||||
|
||||
Returns:
|
||||
JSON Schema字典
|
||||
"""
|
||||
@@ -229,37 +281,39 @@ class MoviePilotToolsManager:
|
||||
|
||||
if "properties" in schema:
|
||||
for field_name, field_info in schema["properties"].items():
|
||||
resolved_field_info = MoviePilotToolsManager._resolve_field_schema(
|
||||
field_info
|
||||
)
|
||||
# 转换字段类型
|
||||
field_type = field_info.get("type", "string")
|
||||
field_description = field_info.get("description", "")
|
||||
field_type = resolved_field_info.get("type", "string")
|
||||
field_description = resolved_field_info.get("description", "")
|
||||
|
||||
# 处理可选字段
|
||||
if field_name not in schema.get("required", []):
|
||||
# 可选字段
|
||||
default_value = field_info.get("default")
|
||||
default_value = resolved_field_info.get("default")
|
||||
properties[field_name] = {
|
||||
"type": field_type,
|
||||
"description": field_description
|
||||
"description": field_description,
|
||||
}
|
||||
if default_value is not None:
|
||||
properties[field_name]["default"] = default_value
|
||||
else:
|
||||
properties[field_name] = {
|
||||
"type": field_type,
|
||||
"description": field_description
|
||||
"description": field_description,
|
||||
}
|
||||
required.append(field_name)
|
||||
|
||||
# 处理枚举类型
|
||||
if "enum" in field_info:
|
||||
properties[field_name]["enum"] = field_info["enum"]
|
||||
if "enum" in resolved_field_info:
|
||||
properties[field_name]["enum"] = resolved_field_info["enum"]
|
||||
|
||||
# 处理数组类型
|
||||
if field_type == "array" and "items" in field_info:
|
||||
properties[field_name]["items"] = field_info["items"]
|
||||
if field_type == "array" and "items" in resolved_field_info:
|
||||
properties[field_name]["items"] = resolved_field_info["items"]
|
||||
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": properties,
|
||||
"required": required
|
||||
}
|
||||
return {"type": "object", "properties": properties, "required": required}
|
||||
|
||||
|
||||
moviepilot_tool_manager = MoviePilotToolsManager()
|
||||
|
||||
@@ -26,11 +26,17 @@ def statistic(name: Optional[str] = None, _: schemas.TokenPayload = Depends(veri
|
||||
if media_statistics:
|
||||
# 汇总各媒体库统计信息
|
||||
ret_statistic = schemas.Statistic()
|
||||
has_episode_count = False
|
||||
for media_statistic in media_statistics:
|
||||
ret_statistic.movie_count += media_statistic.movie_count
|
||||
ret_statistic.tv_count += media_statistic.tv_count
|
||||
ret_statistic.episode_count += media_statistic.episode_count
|
||||
ret_statistic.user_count += media_statistic.user_count
|
||||
ret_statistic.movie_count += media_statistic.movie_count or 0
|
||||
ret_statistic.tv_count += media_statistic.tv_count or 0
|
||||
ret_statistic.user_count += media_statistic.user_count or 0
|
||||
if media_statistic.episode_count is not None:
|
||||
ret_statistic.episode_count += media_statistic.episode_count or 0
|
||||
has_episode_count = True
|
||||
if not has_episode_count:
|
||||
# 所有媒体服务都未提供剧集统计时,返回 None 供前端展示“未获取”。
|
||||
ret_statistic.episode_count = None
|
||||
return ret_statistic
|
||||
else:
|
||||
return schemas.Statistic()
|
||||
|
||||
@@ -5,6 +5,7 @@ from fastapi import APIRouter, Depends, Body
|
||||
from app import schemas
|
||||
from app.chain.download import DownloadChain
|
||||
from app.chain.media import MediaChain
|
||||
from app.core.config import settings
|
||||
from app.core.context import MediaInfo, Context, TorrentInfo
|
||||
from app.core.event import eventmanager
|
||||
from app.core.metainfo import MetaInfo
|
||||
@@ -77,13 +78,14 @@ def add(
|
||||
# 元数据
|
||||
metainfo = MetaInfo(title=torrent_in.title, subtitle=torrent_in.description)
|
||||
# 媒体信息
|
||||
mediainfo = MediaChain().recognize_media(meta=metainfo, tmdbid=tmdbid, doubanid=doubanid)
|
||||
mediainfo = MediaChain().select_recognize_source(
|
||||
log_name=torrent_in.title,
|
||||
log_context=torrent_in.title,
|
||||
native_fn=lambda: MediaChain().recognize_media(meta=metainfo, tmdbid=tmdbid, doubanid=doubanid),
|
||||
plugin_fn=lambda: MediaChain().recognize_help(title=torrent_in.title, org_meta=metainfo)
|
||||
)
|
||||
if not mediainfo:
|
||||
# 尝试使用辅助识别,如果有注册响应事件的话
|
||||
if eventmanager.check(ChainEventType.NameRecognize):
|
||||
mediainfo = MediaChain().recognize_help(title=torrent_in.title, org_meta=metainfo)
|
||||
if not mediainfo:
|
||||
return schemas.Response(success=False, message="无法识别媒体信息")
|
||||
return schemas.Response(success=False, message="无法识别媒体信息")
|
||||
# 种子信息
|
||||
torrentinfo = TorrentInfo()
|
||||
torrentinfo.from_dict(torrent_in.model_dump())
|
||||
|
||||
@@ -4,6 +4,7 @@ import jieba
|
||||
from fastapi import APIRouter, Depends
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import Session
|
||||
from pathlib import Path
|
||||
|
||||
from app import schemas
|
||||
from app.chain.storage import StorageChain
|
||||
@@ -11,7 +12,7 @@ from app.core.event import eventmanager
|
||||
from app.core.security import verify_token
|
||||
from app.db import get_async_db, get_db
|
||||
from app.db.models import User
|
||||
from app.db.models.downloadhistory import DownloadHistory
|
||||
from app.db.models.downloadhistory import DownloadHistory, DownloadFiles
|
||||
from app.db.models.transferhistory import TransferHistory
|
||||
from app.db.user_oper import get_current_active_superuser_async, get_current_active_superuser
|
||||
from app.schemas.types import EventType
|
||||
@@ -98,6 +99,8 @@ def delete_transfer_history(history_in: schemas.TransferHistory,
|
||||
state = StorageChain().delete_media_file(src_fileitem)
|
||||
if not state:
|
||||
return schemas.Response(success=False, message=f"{src_fileitem.path} 删除失败")
|
||||
# 删除下载记录中关联的文件
|
||||
DownloadFiles.delete_by_fullpath(db, Path(src_fileitem.path).as_posix())
|
||||
# 发送事件
|
||||
eventmanager.send_event(
|
||||
EventType.DownloadFileDeleted,
|
||||
|
||||
@@ -32,11 +32,11 @@ def login_access_token(
|
||||
# 如果是需要MFA验证,返回特殊标识
|
||||
if user_or_message == "MFA_REQUIRED":
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
status_code=401,
|
||||
detail="需要双重验证,请提供验证码或使用通行密钥",
|
||||
headers={"X-MFA-Required": "true"}
|
||||
)
|
||||
raise HTTPException(status_code=401, detail=user_or_message)
|
||||
raise HTTPException(status_code=401, detail="用户名或密码错误")
|
||||
|
||||
# 用户等级
|
||||
level = SitesHelper().auth_level
|
||||
|
||||
@@ -1,14 +1,10 @@
|
||||
"""工具API端点
|
||||
通过HTTP API暴露MoviePilot的智能体工具功能
|
||||
"""
|
||||
|
||||
from typing import List, Any, Dict, Annotated, Union
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse, Response
|
||||
|
||||
from app import schemas
|
||||
from app.agent.tools.manager import MoviePilotToolsManager
|
||||
from app.agent.tools.manager import moviepilot_tool_manager
|
||||
from app.core.security import verify_apikey
|
||||
from app.log import logger
|
||||
|
||||
@@ -23,20 +19,29 @@ router = APIRouter()
|
||||
# MCP 协议版本
|
||||
MCP_PROTOCOL_VERSIONS = ["2025-11-25", "2025-06-18", "2024-11-05"]
|
||||
MCP_PROTOCOL_VERSION = MCP_PROTOCOL_VERSIONS[0] # 默认使用最新版本
|
||||
MCP_HIDDEN_TOOLS = {
|
||||
"execute_command",
|
||||
"search_web",
|
||||
"edit_file",
|
||||
"write_file",
|
||||
"read_file",
|
||||
}
|
||||
|
||||
|
||||
def get_tools_manager() -> MoviePilotToolsManager:
|
||||
def list_exposed_tools():
|
||||
"""
|
||||
获取工具管理器实例
|
||||
|
||||
Returns:
|
||||
MoviePilotToolsManager实例
|
||||
获取 MCP 可见工具列表
|
||||
"""
|
||||
return MoviePilotToolsManager()
|
||||
return [
|
||||
tool for tool in moviepilot_tool_manager.list_tools()
|
||||
if tool.name not in MCP_HIDDEN_TOOLS
|
||||
]
|
||||
|
||||
|
||||
def create_jsonrpc_response(request_id: Union[str, int, None], result: Any) -> Dict[str, Any]:
|
||||
"""创建 JSON-RPC 成功响应"""
|
||||
"""
|
||||
创建 JSON-RPC 成功响应
|
||||
"""
|
||||
response = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
@@ -45,8 +50,11 @@ def create_jsonrpc_response(request_id: Union[str, int, None], result: Any) -> D
|
||||
return response
|
||||
|
||||
|
||||
def create_jsonrpc_error(request_id: Union[str, int, None], code: int, message: str, data: Any = None) -> Dict[str, Any]:
|
||||
"""创建 JSON-RPC 错误响应"""
|
||||
def create_jsonrpc_error(request_id: Union[str, int, None], code: int, message: str, data: Any = None) -> Dict[
|
||||
str, Any]:
|
||||
"""
|
||||
创建 JSON-RPC 错误响应
|
||||
"""
|
||||
error = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
@@ -60,8 +68,6 @@ def create_jsonrpc_error(request_id: Union[str, int, None], code: int, message:
|
||||
return error
|
||||
|
||||
|
||||
# ==================== MCP JSON-RPC 端点 ====================
|
||||
|
||||
@router.post("", summary="MCP JSON-RPC 端点", response_model=None)
|
||||
async def mcp_jsonrpc(
|
||||
request: Request,
|
||||
@@ -146,7 +152,9 @@ async def mcp_jsonrpc(
|
||||
|
||||
|
||||
async def handle_initialize(params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""处理初始化请求"""
|
||||
"""
|
||||
处理初始化请求
|
||||
"""
|
||||
protocol_version = params.get("protocolVersion")
|
||||
client_info = params.get("clientInfo", {})
|
||||
|
||||
@@ -161,7 +169,7 @@ async def handle_initialize(params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
else:
|
||||
# 客户端版本不支持,使用服务器默认版本
|
||||
logger.warning(f"协议版本不匹配: 客户端={protocol_version}, 使用服务器版本={negotiated_version}")
|
||||
|
||||
|
||||
return {
|
||||
"protocolVersion": negotiated_version,
|
||||
"capabilities": {
|
||||
@@ -180,9 +188,10 @@ async def handle_initialize(params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
|
||||
|
||||
async def handle_tools_list() -> Dict[str, Any]:
|
||||
"""处理工具列表请求"""
|
||||
manager = get_tools_manager()
|
||||
tools = manager.list_tools()
|
||||
"""
|
||||
处理工具列表请求
|
||||
"""
|
||||
tools = list_exposed_tools()
|
||||
|
||||
# 转换为 MCP 工具格式
|
||||
mcp_tools = []
|
||||
@@ -200,18 +209,21 @@ async def handle_tools_list() -> Dict[str, Any]:
|
||||
|
||||
|
||||
async def handle_tools_call(params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""处理工具调用请求"""
|
||||
"""
|
||||
处理工具调用请求
|
||||
"""
|
||||
tool_name = params.get("name")
|
||||
arguments = params.get("arguments", {})
|
||||
|
||||
if not tool_name:
|
||||
raise ValueError("Missing tool name")
|
||||
|
||||
manager = get_tools_manager()
|
||||
|
||||
try:
|
||||
result_text = await manager.call_tool(tool_name, arguments)
|
||||
|
||||
if tool_name in MCP_HIDDEN_TOOLS:
|
||||
raise ValueError(f"工具 '{tool_name}' 未找到")
|
||||
|
||||
result_text = await moviepilot_tool_manager.call_tool(tool_name, arguments)
|
||||
|
||||
return {
|
||||
"content": [
|
||||
{
|
||||
@@ -243,8 +255,6 @@ async def delete_mcp_session(
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
|
||||
|
||||
# ==================== 兼容的 RESTful API 端点 ====================
|
||||
|
||||
@router.get("/tools", summary="列出所有可用工具", response_model=List[Dict[str, Any]])
|
||||
@@ -257,9 +267,8 @@ async def list_tools(
|
||||
返回每个工具的名称、描述和参数定义
|
||||
"""
|
||||
try:
|
||||
manager = get_tools_manager()
|
||||
# 获取所有工具定义
|
||||
tools = manager.list_tools()
|
||||
tools = list_exposed_tools()
|
||||
|
||||
# 转换为字典格式
|
||||
tools_list = []
|
||||
@@ -289,11 +298,10 @@ async def call_tool(
|
||||
工具执行结果
|
||||
"""
|
||||
try:
|
||||
# 使用当前用户ID创建管理器实例
|
||||
manager = get_tools_manager()
|
||||
if request.tool_name in MCP_HIDDEN_TOOLS:
|
||||
raise ValueError(f"工具 '{request.tool_name}' 未找到")
|
||||
|
||||
# 调用工具
|
||||
result_text = await manager.call_tool(request.tool_name, request.arguments)
|
||||
result_text = await moviepilot_tool_manager.call_tool(request.tool_name, request.arguments)
|
||||
|
||||
return schemas.ToolCallResponse(
|
||||
success=True,
|
||||
@@ -319,9 +327,8 @@ async def get_tool_info(
|
||||
工具的详细信息,包括名称、描述和参数定义
|
||||
"""
|
||||
try:
|
||||
manager = get_tools_manager()
|
||||
# 获取所有工具
|
||||
tools = manager.list_tools()
|
||||
tools = list_exposed_tools()
|
||||
|
||||
# 查找指定工具
|
||||
for tool in tools:
|
||||
@@ -352,9 +359,8 @@ async def get_tool_schema(
|
||||
工具的JSON Schema定义
|
||||
"""
|
||||
try:
|
||||
manager = get_tools_manager()
|
||||
# 获取所有工具
|
||||
tools = manager.list_tools()
|
||||
tools = list_exposed_tools()
|
||||
|
||||
# 查找指定工具
|
||||
for tool in tools:
|
||||
|
||||
@@ -11,7 +11,10 @@ from app.core.context import Context
|
||||
from app.core.event import eventmanager
|
||||
from app.core.metainfo import MetaInfo, MetaInfoPath
|
||||
from app.core.security import verify_token, verify_apitoken
|
||||
from app.db.models import User
|
||||
from app.db.user_oper import get_current_active_user, get_current_active_superuser
|
||||
from app.schemas import MediaType, MediaRecognizeConvertEventData
|
||||
from app.schemas.category import CategoryConfig
|
||||
from app.schemas.types import ChainEventType
|
||||
|
||||
router = APIRouter()
|
||||
@@ -131,6 +134,26 @@ def scrape(fileitem: schemas.FileItem,
|
||||
return schemas.Response(success=True, message=f"{fileitem.path} 刮削完成")
|
||||
|
||||
|
||||
@router.get("/category/config", summary="获取分类策略配置", response_model=schemas.Response)
|
||||
def get_category_config(_: User = Depends(get_current_active_user)):
|
||||
"""
|
||||
获取分类策略配置
|
||||
"""
|
||||
config = MediaChain().category_config()
|
||||
return schemas.Response(success=True, data=config.model_dump())
|
||||
|
||||
|
||||
@router.post("/category/config", summary="保存分类策略配置", response_model=schemas.Response)
|
||||
def save_category_config(config: CategoryConfig, _: User = Depends(get_current_active_superuser)):
|
||||
"""
|
||||
保存分类策略配置
|
||||
"""
|
||||
if MediaChain().save_category_config(config):
|
||||
return schemas.Response(success=True, message="保存成功")
|
||||
else:
|
||||
return schemas.Response(success=False, message="保存失败")
|
||||
|
||||
|
||||
@router.get("/category", summary="查询自动分类配置", response_model=dict)
|
||||
async def category(_: schemas.TokenPayload = Depends(verify_token)) -> Any:
|
||||
"""
|
||||
@@ -172,7 +195,7 @@ async def seasons(mediaid: Optional[str] = None,
|
||||
tmdbid = int(mediaid[5:])
|
||||
seasons_info = await TmdbChain().async_tmdb_seasons(tmdbid=tmdbid)
|
||||
if seasons_info:
|
||||
if season:
|
||||
if season is not None:
|
||||
return [sea for sea in seasons_info if sea.season_number == season]
|
||||
return seasons_info
|
||||
if title:
|
||||
@@ -184,11 +207,11 @@ async def seasons(mediaid: Optional[str] = None,
|
||||
if settings.RECOGNIZE_SOURCE == "themoviedb":
|
||||
seasons_info = await TmdbChain().async_tmdb_seasons(tmdbid=mediainfo.tmdb_id)
|
||||
if seasons_info:
|
||||
if season:
|
||||
if season is not None:
|
||||
return [sea for sea in seasons_info if sea.season_number == season]
|
||||
return seasons_info
|
||||
else:
|
||||
sea = season or 1
|
||||
sea = season if season is not None else 1
|
||||
return [schemas.MediaSeason(
|
||||
season_number=sea,
|
||||
poster_path=mediainfo.poster_path,
|
||||
|
||||
@@ -54,7 +54,7 @@ async def exists_local(title: Optional[str] = None,
|
||||
判断本地是否存在
|
||||
"""
|
||||
meta = MetaInfo(title)
|
||||
if not season:
|
||||
if season is None:
|
||||
season = meta.begin_season
|
||||
# 返回对象
|
||||
ret_info = {}
|
||||
@@ -82,8 +82,8 @@ def exists(media_in: schemas.MediaInfo,
|
||||
mediainfo.from_dict(media_in.model_dump())
|
||||
existsinfo: schemas.ExistMediaInfo = MediaServerChain().media_exists(mediainfo=mediainfo)
|
||||
if not existsinfo:
|
||||
return []
|
||||
if media_in.season:
|
||||
return {}
|
||||
if media_in.season is not None:
|
||||
return {
|
||||
media_in.season: existsinfo.seasons.get(media_in.season) or []
|
||||
}
|
||||
@@ -101,7 +101,7 @@ def not_exists(media_in: schemas.MediaInfo,
|
||||
mtype = MediaType(media_in.type) if media_in.type else None
|
||||
if mtype:
|
||||
meta.type = mtype
|
||||
if media_in.season:
|
||||
if media_in.season is not None:
|
||||
meta.begin_season = media_in.season
|
||||
meta.type = MediaType.TV
|
||||
if media_in.year:
|
||||
|
||||
@@ -86,7 +86,10 @@ def wechat_verify(echostr: str, msg_signature: str, timestamp: Union[str, int],
|
||||
if not client_configs:
|
||||
return "未找到对应的消息配置"
|
||||
client_config = next((config for config in client_configs if
|
||||
config.type == "wechat" and config.enabled and (not source or config.name == source)), None)
|
||||
config.type == "wechat"
|
||||
and config.enabled
|
||||
and config.config.get("WECHAT_MODE", "app") != "bot"
|
||||
and (not source or config.name == source)), None)
|
||||
if not client_config:
|
||||
return "未找到对应的消息配置"
|
||||
try:
|
||||
|
||||
@@ -24,6 +24,75 @@ from app.utils.otp import OtpUtils
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# ==================== 辅助函数 ====================
|
||||
|
||||
def _build_credential_list(passkeys: list[PassKey]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
构建凭证列表
|
||||
|
||||
:param passkeys: PassKey 列表
|
||||
:return: 凭证字典列表
|
||||
"""
|
||||
return [
|
||||
{
|
||||
'credential_id': pk.credential_id,
|
||||
'transports': pk.transports
|
||||
}
|
||||
for pk in passkeys
|
||||
] if passkeys else []
|
||||
|
||||
|
||||
def _extract_and_standardize_credential_id(credential: dict) -> str:
|
||||
"""
|
||||
从凭证中提取并标准化 credential_id
|
||||
|
||||
:param credential: 凭证字典
|
||||
:return: 标准化后的 credential_id
|
||||
:raises ValueError: 如果凭证无效
|
||||
"""
|
||||
credential_id_raw = credential.get('id') or credential.get('rawId')
|
||||
if not credential_id_raw:
|
||||
raise ValueError("无效的凭证")
|
||||
return PassKeyHelper.standardize_credential_id(credential_id_raw)
|
||||
|
||||
|
||||
def _verify_passkey_and_update(
|
||||
credential: dict,
|
||||
challenge: str,
|
||||
passkey: PassKey
|
||||
) -> tuple[bool, int]:
|
||||
"""
|
||||
验证 PassKey 并更新使用时间和签名计数
|
||||
|
||||
:param credential: 凭证字典
|
||||
:param challenge: 挑战值
|
||||
:param passkey: PassKey 对象
|
||||
:return: (验证是否成功, 新的签名计数)
|
||||
"""
|
||||
success, new_sign_count = PassKeyHelper.verify_authentication_response(
|
||||
credential=credential,
|
||||
expected_challenge=challenge,
|
||||
credential_public_key=passkey.public_key,
|
||||
credential_current_sign_count=passkey.sign_count
|
||||
)
|
||||
|
||||
if success:
|
||||
passkey.update_last_used(db=None, sign_count=new_sign_count)
|
||||
|
||||
return success, new_sign_count
|
||||
|
||||
|
||||
async def _check_user_has_passkey(db: AsyncSession, user_id: int) -> bool:
|
||||
"""
|
||||
检查用户是否有 PassKey
|
||||
|
||||
:param db: 数据库会话
|
||||
:param user_id: 用户 ID
|
||||
:return: 是否有 PassKey
|
||||
"""
|
||||
return bool(await PassKey.async_get_by_user_id(db=db, user_id=user_id))
|
||||
|
||||
|
||||
# ==================== 请求模型 ====================
|
||||
|
||||
class OtpVerifyRequest(schemas.BaseModel):
|
||||
@@ -55,7 +124,7 @@ async def mfa_status(username: str, db: AsyncSession = Depends(get_async_db)) ->
|
||||
has_otp = user.is_otp
|
||||
|
||||
# 检查是否有PassKey
|
||||
has_passkey = bool(await PassKey.async_get_by_user_id(db=db, user_id=user.id))
|
||||
has_passkey = await _check_user_has_passkey(db, user.id)
|
||||
|
||||
# 只要有任何一种验证方式,就需要双重验证
|
||||
return schemas.Response(success=(has_otp or has_passkey))
|
||||
@@ -92,9 +161,9 @@ async def otp_disable(
|
||||
current_user: User = Depends(get_current_active_user_async)
|
||||
) -> Any:
|
||||
"""关闭当前用户的 OTP 验证功能"""
|
||||
# 安全检查:如果存在 PassKey,不允许关闭 OTP
|
||||
has_passkey = bool(await PassKey.async_get_by_user_id(db=db, user_id=current_user.id))
|
||||
if has_passkey:
|
||||
# 安全检查:如果存在 PassKey,默认不允许关闭 OTP,除非配置允许
|
||||
has_passkey = await _check_user_has_passkey(db, current_user.id)
|
||||
if has_passkey and not settings.PASSKEY_ALLOW_REGISTER_WITHOUT_OTP:
|
||||
return schemas.Response(
|
||||
success=False,
|
||||
message="您已注册通行密钥,为了防止域名配置变更导致无法登录,请先删除所有通行密钥再关闭 OTP 验证"
|
||||
@@ -138,8 +207,8 @@ def passkey_register_start(
|
||||
) -> Any:
|
||||
"""开始注册 PassKey - 生成注册选项"""
|
||||
try:
|
||||
# 安全检查:必须先启用 OTP
|
||||
if not current_user.is_otp:
|
||||
# 安全检查:默认需要先启用 OTP,除非配置允许在未启用 OTP 时注册
|
||||
if not current_user.is_otp and not settings.PASSKEY_ALLOW_REGISTER_WITHOUT_OTP:
|
||||
return schemas.Response(
|
||||
success=False,
|
||||
message="为了确保在域名配置错误时仍能找回访问权限,请先启用 OTP 验证码再注册通行密钥"
|
||||
@@ -147,13 +216,7 @@ def passkey_register_start(
|
||||
|
||||
# 获取用户已有的PassKey
|
||||
existing_passkeys = PassKey.get_by_user_id(db=None, user_id=current_user.id)
|
||||
existing_credentials = [
|
||||
{
|
||||
'credential_id': pk.credential_id,
|
||||
'transports': pk.transports
|
||||
}
|
||||
for pk in existing_passkeys
|
||||
] if existing_passkeys else None
|
||||
existing_credentials = _build_credential_list(existing_passkeys) if existing_passkeys else None
|
||||
|
||||
# 生成注册选项
|
||||
options_json, challenge = PassKeyHelper.generate_registration_options(
|
||||
@@ -233,26 +296,15 @@ def passkey_authenticate_start(
|
||||
# 如果指定了用户名,只允许该用户的PassKey
|
||||
if passkey_req.username:
|
||||
user = User.get_by_name(db=None, name=passkey_req.username)
|
||||
if not user:
|
||||
existing_passkeys = PassKey.get_by_user_id(db=None, user_id=user.id) if user else None
|
||||
|
||||
if not user or not existing_passkeys:
|
||||
return schemas.Response(
|
||||
success=False,
|
||||
message="用户不存在"
|
||||
message="认证失败"
|
||||
)
|
||||
|
||||
existing_passkeys = PassKey.get_by_user_id(db=None, user_id=user.id)
|
||||
if not existing_passkeys:
|
||||
return schemas.Response(
|
||||
success=False,
|
||||
message="该用户未注册通行密钥"
|
||||
)
|
||||
|
||||
existing_credentials = [
|
||||
{
|
||||
'credential_id': pk.credential_id,
|
||||
'transports': pk.transports
|
||||
}
|
||||
for pk in existing_passkeys
|
||||
]
|
||||
|
||||
existing_credentials = _build_credential_list(existing_passkeys)
|
||||
|
||||
# 生成认证选项
|
||||
options_json, challenge = PassKeyHelper.generate_authentication_options(
|
||||
@@ -270,7 +322,7 @@ def passkey_authenticate_start(
|
||||
logger.error(f"生成PassKey认证选项失败: {e}")
|
||||
return schemas.Response(
|
||||
success=False,
|
||||
message=f"生成认证选项失败: {str(e)}"
|
||||
message="认证失败"
|
||||
)
|
||||
|
||||
|
||||
@@ -280,37 +332,28 @@ def passkey_authenticate_finish(
|
||||
) -> Any:
|
||||
"""完成 PassKey 认证 - 验证凭证并返回 token"""
|
||||
try:
|
||||
# 从credential中提取credential_id
|
||||
credential_id_raw = passkey_req.credential.get('id') or passkey_req.credential.get('rawId')
|
||||
if not credential_id_raw:
|
||||
raise HTTPException(status_code=400, detail="无效的凭证")
|
||||
# 提取并标准化凭证ID
|
||||
try:
|
||||
credential_id = _extract_and_standardize_credential_id(passkey_req.credential)
|
||||
except ValueError as e:
|
||||
logger.warning(f"PassKey认证失败,提供的凭证无效: {e}")
|
||||
raise HTTPException(status_code=401, detail="认证失败")
|
||||
|
||||
# 标准化凭证ID
|
||||
credential_id = PassKeyHelper.standardize_credential_id(credential_id_raw)
|
||||
|
||||
# 查找PassKey
|
||||
# 查找PassKey并获取用户
|
||||
passkey = PassKey.get_by_credential_id(db=None, credential_id=credential_id)
|
||||
if not passkey:
|
||||
raise HTTPException(status_code=401, detail="通行密钥不存在或已失效")
|
||||
user = User.get_by_id(db=None, user_id=passkey.user_id) if passkey else None
|
||||
if not passkey or not user or not user.is_active:
|
||||
raise HTTPException(status_code=401, detail="认证失败")
|
||||
|
||||
# 获取用户
|
||||
user = User.get_by_id(db=None, user_id=passkey.user_id)
|
||||
if not user or not user.is_active:
|
||||
raise HTTPException(status_code=401, detail="用户不存在或已禁用")
|
||||
|
||||
# 验证认证响应
|
||||
success, new_sign_count = PassKeyHelper.verify_authentication_response(
|
||||
# 验证认证响应并更新
|
||||
success, _ = _verify_passkey_and_update(
|
||||
credential=passkey_req.credential,
|
||||
expected_challenge=passkey_req.challenge,
|
||||
credential_public_key=passkey.public_key,
|
||||
credential_current_sign_count=passkey.sign_count
|
||||
challenge=passkey_req.challenge,
|
||||
passkey=passkey
|
||||
)
|
||||
|
||||
if not success:
|
||||
raise HTTPException(status_code=401, detail="通行密钥验证失败")
|
||||
|
||||
# 更新使用时间和签名计数
|
||||
passkey.update_last_used(db=None, sign_count=new_sign_count)
|
||||
raise HTTPException(status_code=401, detail="认证失败")
|
||||
|
||||
logger.info(f"用户 {user.name} 通过PassKey认证成功")
|
||||
|
||||
@@ -339,7 +382,7 @@ def passkey_authenticate_finish(
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"PassKey认证失败: {e}")
|
||||
raise HTTPException(status_code=401, detail=f"认证失败: {str(e)}")
|
||||
raise HTTPException(status_code=401, detail="认证失败")
|
||||
|
||||
|
||||
@router.get("/passkey/list", summary="获取当前用户的 PassKey 列表", response_model=schemas.Response)
|
||||
@@ -413,16 +456,12 @@ def passkey_verify_mfa(
|
||||
) -> Any:
|
||||
"""使用 PassKey 进行二次验证(MFA)"""
|
||||
try:
|
||||
# 从credential中提取credential_id
|
||||
credential_id_raw = passkey_req.credential.get('id') or passkey_req.credential.get('rawId')
|
||||
if not credential_id_raw:
|
||||
return schemas.Response(
|
||||
success=False,
|
||||
message="无效的凭证"
|
||||
)
|
||||
|
||||
# 标准化凭证ID
|
||||
credential_id = PassKeyHelper.standardize_credential_id(credential_id_raw)
|
||||
# 提取并标准化凭证ID
|
||||
try:
|
||||
credential_id = _extract_and_standardize_credential_id(passkey_req.credential)
|
||||
except ValueError as e:
|
||||
logger.warning(f"PassKey二次验证失败,提供的凭证无效: {e}")
|
||||
return schemas.Response(success=False, message="验证失败")
|
||||
|
||||
# 查找PassKey(必须属于当前用户)
|
||||
passkey = PassKey.get_by_credential_id(db=None, credential_id=credential_id)
|
||||
@@ -432,12 +471,11 @@ def passkey_verify_mfa(
|
||||
message="通行密钥不存在或不属于当前用户"
|
||||
)
|
||||
|
||||
# 验证认证响应
|
||||
success, new_sign_count = PassKeyHelper.verify_authentication_response(
|
||||
# 验证认证响应并更新
|
||||
success, _ = _verify_passkey_and_update(
|
||||
credential=passkey_req.credential,
|
||||
expected_challenge=passkey_req.challenge,
|
||||
credential_public_key=passkey.public_key,
|
||||
credential_current_sign_count=passkey.sign_count
|
||||
challenge=passkey_req.challenge,
|
||||
passkey=passkey
|
||||
)
|
||||
|
||||
if not success:
|
||||
@@ -446,9 +484,6 @@ def passkey_verify_mfa(
|
||||
message="通行密钥验证失败"
|
||||
)
|
||||
|
||||
# 更新使用时间和签名计数
|
||||
passkey.update_last_used(db=None, sign_count=new_sign_count)
|
||||
|
||||
logger.info(f"用户 {current_user.name} 通过PassKey二次验证成功")
|
||||
|
||||
return schemas.Response(
|
||||
@@ -459,5 +494,5 @@ def passkey_verify_mfa(
|
||||
logger.error(f"PassKey二次验证失败: {e}")
|
||||
return schemas.Response(
|
||||
success=False,
|
||||
message=f"验证失败: {str(e)}"
|
||||
message="验证失败"
|
||||
)
|
||||
|
||||
@@ -360,7 +360,18 @@ async def plugin_static_file(plugin_id: str, filepath: str):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Forbidden")
|
||||
|
||||
plugin_base_dir = AsyncPath(settings.ROOT_PATH) / "app" / "plugins" / plugin_id.lower()
|
||||
plugin_file_path = plugin_base_dir / filepath
|
||||
plugin_file_path = plugin_base_dir / filepath.lstrip('/')
|
||||
|
||||
try:
|
||||
resolved_base = await plugin_base_dir.resolve()
|
||||
resolved_file = await plugin_file_path.resolve()
|
||||
except Exception:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid path")
|
||||
|
||||
if not resolved_file.is_relative_to(resolved_base):
|
||||
logger.warning(f"Static File API: Path traversal attempt detected: {plugin_id}/{filepath}")
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Forbidden")
|
||||
|
||||
if not await plugin_file_path.exists():
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"{plugin_file_path} 不存在")
|
||||
if not await plugin_file_path.is_file():
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
from typing import List, Any, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
from fastapi import APIRouter, Depends, Body
|
||||
|
||||
from app import schemas
|
||||
from app.chain.media import MediaChain
|
||||
from app.chain.search import SearchChain
|
||||
from app.chain.ai_recommend import AIRecommendChain
|
||||
from app.core.config import settings
|
||||
from app.core.event import eventmanager
|
||||
from app.core.metainfo import MetaInfo
|
||||
from app.core.security import verify_token
|
||||
from app.log import logger
|
||||
from app.schemas import MediaRecognizeConvertEventData
|
||||
from app.schemas.types import MediaType, ChainEventType
|
||||
|
||||
@@ -36,6 +38,9 @@ async def search_by_id(mediaid: str,
|
||||
"""
|
||||
根据TMDBID/豆瓣ID精确搜索站点资源 tmdb:/douban:/bangumi:
|
||||
"""
|
||||
# 取消正在运行的AI推荐(会清除数据库缓存)
|
||||
AIRecommendChain().cancel_ai_recommend()
|
||||
|
||||
if mtype:
|
||||
media_type = MediaType(mtype)
|
||||
else:
|
||||
@@ -159,6 +164,9 @@ async def search_by_title(keyword: Optional[str] = None,
|
||||
"""
|
||||
根据名称模糊搜索站点资源,支持分页,关键词为空是返回首页资源
|
||||
"""
|
||||
# 取消正在运行的AI推荐并清除数据库缓存
|
||||
AIRecommendChain().cancel_ai_recommend()
|
||||
|
||||
torrents = await SearchChain().async_search_by_title(
|
||||
title=keyword, page=page,
|
||||
sites=[int(site) for site in sites.split(",") if site] if sites else None,
|
||||
@@ -167,3 +175,87 @@ async def search_by_title(keyword: Optional[str] = None,
|
||||
if not torrents:
|
||||
return schemas.Response(success=False, message="未搜索到任何资源")
|
||||
return schemas.Response(success=True, data=[torrent.to_dict() for torrent in torrents])
|
||||
|
||||
|
||||
@router.post("/recommend", summary="AI推荐资源", response_model=schemas.Response)
|
||||
async def recommend_search_results(
|
||||
filtered_indices: Optional[List[int]] = Body(None, embed=True, description="筛选后的索引列表"),
|
||||
check_only: bool = Body(False, embed=True, description="仅检查状态,不启动新任务"),
|
||||
force: bool = Body(False, embed=True, description="强制重新推荐,清除旧结果"),
|
||||
_: schemas.TokenPayload = Depends(verify_token)) -> Any:
|
||||
"""
|
||||
AI推荐资源 - 轮询接口
|
||||
前端轮询此接口,发送筛选后的索引(如果有筛选)
|
||||
后端根据请求变化自动取消旧任务并启动新任务
|
||||
|
||||
参数:
|
||||
- filtered_indices: 筛选后的索引列表(可选,为空或不提供时使用所有结果)
|
||||
- check_only: 仅检查状态(首次打开页面时使用,避免触发不必要的重新推理)
|
||||
- force: 强制重新推荐(清除旧结果并重新启动)
|
||||
|
||||
返回数据结构:
|
||||
{
|
||||
"success": bool,
|
||||
"message": string, // 错误信息(仅在错误时存在)
|
||||
"data": {
|
||||
"status": string, // 状态: disabled | idle | running | completed | error
|
||||
"results": array // 推荐结果(仅status=completed时存在)
|
||||
}
|
||||
}
|
||||
"""
|
||||
# 从缓存获取上次搜索结果
|
||||
results = await SearchChain().async_last_search_results() or []
|
||||
if not results:
|
||||
return schemas.Response(success=False, message="没有可用的搜索结果", data={
|
||||
"status": "error"
|
||||
})
|
||||
|
||||
recommend_chain = AIRecommendChain()
|
||||
|
||||
# 如果是强制模式,先取消并清除旧结果,然后直接启动新任务
|
||||
if force:
|
||||
# 检查功能是否启用
|
||||
if not settings.AI_AGENT_ENABLE or not settings.AI_RECOMMEND_ENABLED:
|
||||
return schemas.Response(success=True, data={
|
||||
"status": "disabled"
|
||||
})
|
||||
logger.info("收到新推荐请求,清除旧结果并启动新任务")
|
||||
recommend_chain.cancel_ai_recommend()
|
||||
recommend_chain.start_recommend_task(filtered_indices, len(results), results)
|
||||
# 直接返回运行中状态
|
||||
return schemas.Response(success=True, data={
|
||||
"status": "running"
|
||||
})
|
||||
|
||||
# 如果是仅检查模式,不传递 filtered_indices(避免触发请求变化检测)
|
||||
if check_only:
|
||||
# 返回当前运行状态,不做任何任务启动或取消操作
|
||||
current_status = recommend_chain.get_current_status_only()
|
||||
# 如果有错误,将错误信息放到message中
|
||||
if current_status.get("status") == "error":
|
||||
error_msg = current_status.pop("error", "未知错误")
|
||||
return schemas.Response(success=False, message=error_msg, data=current_status)
|
||||
return schemas.Response(success=True, data=current_status)
|
||||
|
||||
# 获取当前状态(会检测请求是否变化)
|
||||
status_data = recommend_chain.get_status(filtered_indices, len(results))
|
||||
|
||||
# 如果功能未启用,直接返回禁用状态
|
||||
if status_data.get("status") == "disabled":
|
||||
return schemas.Response(success=True, data=status_data)
|
||||
|
||||
# 如果是空闲状态,启动新任务
|
||||
if status_data["status"] == "idle":
|
||||
recommend_chain.start_recommend_task(filtered_indices, len(results), results)
|
||||
# 立即返回运行中状态
|
||||
return schemas.Response(success=True, data={
|
||||
"status": "running"
|
||||
})
|
||||
|
||||
# 如果有错误,将错误信息放到message中
|
||||
if status_data.get("status") == "error":
|
||||
error_msg = status_data.pop("error", "未知错误")
|
||||
return schemas.Response(success=False, message=error_msg, data=status_data)
|
||||
|
||||
# 返回当前状态
|
||||
return schemas.Response(success=True, data=status_data)
|
||||
|
||||
@@ -92,10 +92,14 @@ async def update_site(
|
||||
# 校正地址格式
|
||||
_scheme, _netloc = StringUtils.get_url_netloc(site_in.url)
|
||||
site_in.url = f"{_scheme}://{_netloc}/"
|
||||
site_in.domain = StringUtils.get_url_domain(site_in.url)
|
||||
await site.async_update(db, site_in.model_dump())
|
||||
# 通知站点更新
|
||||
await eventmanager.async_send_event(EventType.SiteUpdated, {
|
||||
"domain": site_in.domain
|
||||
"site_id": site_in.id,
|
||||
"domain": site_in.domain,
|
||||
"name": site_in.name,
|
||||
"site_url": site_in.url
|
||||
})
|
||||
return schemas.Response(success=True)
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from datetime import datetime
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Any, List, Optional
|
||||
|
||||
@@ -31,6 +31,17 @@ def qrcode(name: str, _: schemas.TokenPayload = Depends(verify_token)) -> Any:
|
||||
return schemas.Response(success=False, message=errmsg)
|
||||
|
||||
|
||||
@router.get("/auth_url/{name}", summary="获取 OAuth2 授权 URL", response_model=schemas.Response)
|
||||
def auth_url(name: str, _: schemas.TokenPayload = Depends(verify_token)) -> Any:
|
||||
"""
|
||||
获取 OAuth2 授权 URL
|
||||
"""
|
||||
auth_data, errmsg = StorageChain().generate_auth_url(name)
|
||||
if auth_data:
|
||||
return schemas.Response(success=True, data=auth_data)
|
||||
return schemas.Response(success=False, message=errmsg)
|
||||
|
||||
|
||||
@router.get("/check/{name}", summary="二维码登录确认", response_model=schemas.Response)
|
||||
def check(name: str, ck: Optional[str] = None, t: Optional[str] = None,
|
||||
_: schemas.TokenPayload = Depends(verify_token)) -> Any:
|
||||
@@ -83,7 +94,7 @@ def list_files(fileitem: schemas.FileItem,
|
||||
if sort == "name":
|
||||
file_list.sort(key=lambda x: StringUtils.natural_sort_key(x.name or ""))
|
||||
else:
|
||||
file_list.sort(key=lambda x: x.modify_time or datetime.min, reverse=True)
|
||||
file_list.sort(key=lambda x: x.modify_time or -math.inf, reverse=True)
|
||||
return file_list
|
||||
|
||||
|
||||
@@ -167,7 +178,7 @@ def rename(fileitem: schemas.FileItem,
|
||||
# 重命名目录内文件
|
||||
if recursive:
|
||||
transferchain = TransferChain()
|
||||
media_exts = settings.RMT_MEDIAEXT + settings.RMT_SUBEXT + settings.RMT_AUDIO_TRACK_EXT
|
||||
media_exts = settings.RMT_MEDIAEXT + settings.RMT_SUBEXT + settings.RMT_AUDIOEXT
|
||||
# 递归修改目录内文件(智能识别命名)
|
||||
sub_files: List[schemas.FileItem] = StorageChain().list_files(fileitem)
|
||||
if sub_files:
|
||||
|
||||
@@ -199,7 +199,7 @@ async def subscribe_mediaid(
|
||||
# 使用名称检查订阅
|
||||
if title_check and title:
|
||||
meta = MetaInfo(title)
|
||||
if season:
|
||||
if season is not None:
|
||||
meta.begin_season = season
|
||||
result = await Subscribe.async_get_by_title(db, title=meta.name, season=meta.begin_season)
|
||||
|
||||
|
||||
@@ -130,26 +130,53 @@ async def cache_img(
|
||||
def get_global_setting(token: str):
|
||||
"""
|
||||
查询非敏感系统设置(默认鉴权)
|
||||
仅包含登录前UI初始化必需的字段
|
||||
"""
|
||||
if token != "moviepilot":
|
||||
raise HTTPException(status_code=403, detail="Forbidden")
|
||||
|
||||
# 白名单模式,仅包含前端业务逻辑必需的字段
|
||||
# 白名单模式,仅包含登录前UI初始化必需的字段
|
||||
info = settings.model_dump(
|
||||
include={
|
||||
"TMDB_IMAGE_DOMAIN",
|
||||
"GLOBAL_IMAGE_CACHE",
|
||||
"ADVANCED_MODE",
|
||||
"RECOGNIZE_SOURCE",
|
||||
"SEARCH_SOURCE"
|
||||
}
|
||||
)
|
||||
# 追加版本信息(用于版本检查)
|
||||
info.update({
|
||||
"FRONTEND_VERSION": SystemChain.get_frontend_version(),
|
||||
"BACKEND_VERSION": APP_VERSION
|
||||
})
|
||||
return schemas.Response(success=True,
|
||||
data=info)
|
||||
|
||||
|
||||
@router.get("/global/user", summary="查询用户相关系统设置", response_model=schemas.Response)
|
||||
async def get_user_global_setting(_: User = Depends(get_current_active_user_async)):
|
||||
"""
|
||||
查询用户相关系统设置(登录后获取)
|
||||
包含业务功能相关的配置和用户权限信息
|
||||
"""
|
||||
# 业务功能相关的配置字段
|
||||
info = settings.model_dump(
|
||||
include={
|
||||
"RECOGNIZE_SOURCE",
|
||||
"SEARCH_SOURCE",
|
||||
"AI_RECOMMEND_ENABLED",
|
||||
"PASSKEY_ALLOW_REGISTER_WITHOUT_OTP"
|
||||
}
|
||||
)
|
||||
# 智能助手总开关未开启,智能推荐状态强制返回False
|
||||
if not settings.AI_AGENT_ENABLE:
|
||||
info["AI_RECOMMEND_ENABLED"] = False
|
||||
|
||||
# 追加用户唯一ID和订阅分享管理权限
|
||||
share_admin = SubscribeHelper().is_admin_user()
|
||||
info.update({
|
||||
"USER_UNIQUE_ID": SubscribeHelper().get_user_uuid(),
|
||||
"SUBSCRIBE_SHARE_MANAGE": share_admin,
|
||||
"WORKFLOW_SHARE_MANAGE": share_admin
|
||||
"WORKFLOW_SHARE_MANAGE": share_admin,
|
||||
})
|
||||
return schemas.Response(success=True,
|
||||
data=info)
|
||||
@@ -588,7 +615,10 @@ def run_scheduler(jobid: str,
|
||||
"""
|
||||
if not jobid:
|
||||
return schemas.Response(success=False, message="命令不能为空!")
|
||||
Scheduler().start(jobid)
|
||||
if jobid in {"recommend_refresh", "cookiecloud"}:
|
||||
Scheduler().start(jobid, manual=True)
|
||||
else:
|
||||
Scheduler().start(jobid)
|
||||
return schemas.Response(success=True)
|
||||
|
||||
|
||||
@@ -601,5 +631,8 @@ def run_scheduler2(jobid: str,
|
||||
if not jobid:
|
||||
return schemas.Response(success=False, message="命令不能为空!")
|
||||
|
||||
Scheduler().start(jobid)
|
||||
if jobid in {"recommend_refresh", "cookiecloud"}:
|
||||
Scheduler().start(jobid, manual=True)
|
||||
else:
|
||||
Scheduler().start(jobid)
|
||||
return schemas.Response(success=True)
|
||||
|
||||
@@ -93,6 +93,8 @@ def manual_transfer(transer_item: ManualTransferItem,
|
||||
:param _: Token校验
|
||||
"""
|
||||
force = False
|
||||
downloader = None
|
||||
download_hash = None
|
||||
target_path = Path(transer_item.target_path) if transer_item.target_path else None
|
||||
if transer_item.logid:
|
||||
# 查询历史记录
|
||||
@@ -101,6 +103,8 @@ def manual_transfer(transer_item: ManualTransferItem,
|
||||
return schemas.Response(success=False, message=f"整理记录不存在,ID:{transer_item.logid}")
|
||||
# 强制转移
|
||||
force = True
|
||||
downloader = history.downloader
|
||||
download_hash = history.download_hash
|
||||
if history.status and ("move" in history.mode):
|
||||
# 重新整理成功的转移,则使用成功的 dest 做 in_path
|
||||
src_fileitem = FileItem(**history.dest_fileitem)
|
||||
@@ -121,6 +125,7 @@ def manual_transfer(transer_item: ManualTransferItem,
|
||||
transer_item.tmdbid = int(history.tmdbid) if history.tmdbid else transer_item.tmdbid
|
||||
transer_item.doubanid = str(history.doubanid) if history.doubanid else transer_item.doubanid
|
||||
transer_item.season = int(str(history.seasons).replace("S", "")) if history.seasons else transer_item.season
|
||||
transer_item.episode_group = history.episode_group or transer_item.episode_group
|
||||
if history.episodes:
|
||||
if "-" in str(history.episodes):
|
||||
# E01-E03多集合并
|
||||
@@ -138,8 +143,14 @@ def manual_transfer(transer_item: ManualTransferItem,
|
||||
else:
|
||||
return schemas.Response(success=False, message=f"缺少参数")
|
||||
|
||||
# 类型
|
||||
mtype = MediaType(transer_item.type_name) if transer_item.type_name else None
|
||||
# 类型(“自动/auto/none”按未指定处理)
|
||||
mtype = None
|
||||
type_name = str(transer_item.type_name).strip() if transer_item.type_name else ""
|
||||
if type_name and type_name.lower() not in {"自动", "auto", "none"}:
|
||||
try:
|
||||
mtype = MediaType(type_name)
|
||||
except ValueError:
|
||||
return schemas.Response(success=False, message=f"不支持的媒体类型:{type_name}")
|
||||
# 自定义格式
|
||||
epformat = None
|
||||
if transer_item.episode_offset or transer_item.episode_part \
|
||||
@@ -167,7 +178,9 @@ def manual_transfer(transer_item: ManualTransferItem,
|
||||
library_type_folder=transer_item.library_type_folder,
|
||||
library_category_folder=transer_item.library_category_folder,
|
||||
force=force,
|
||||
background=background
|
||||
background=background,
|
||||
downloader=downloader,
|
||||
download_hash=download_hash
|
||||
)
|
||||
# 失败
|
||||
if not state:
|
||||
|
||||
@@ -4,6 +4,7 @@ import pickle
|
||||
import traceback
|
||||
from abc import ABCMeta
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional, Any, Tuple, List, Set, Union, Dict
|
||||
|
||||
@@ -25,6 +26,7 @@ from app.helper.service import ServiceConfigHelper
|
||||
from app.log import logger
|
||||
from app.schemas import TransferInfo, TransferTorrent, ExistMediaInfo, DownloadingTorrent, CommingMessage, Notification, \
|
||||
WebhookEventInfo, TmdbEpisode, MediaPerson, FileItem, TransferDirectoryConf
|
||||
from app.schemas.category import CategoryConfig
|
||||
from app.schemas.types import TorrentStatus, MediaType, MediaImageType, EventType, MessageChannel
|
||||
from app.utils.object import ObjectUtils
|
||||
|
||||
@@ -250,6 +252,7 @@ class ChainBase(metaclass=ABCMeta):
|
||||
# 中止继续执行
|
||||
break
|
||||
except Exception as err:
|
||||
logger.error(traceback.format_exc())
|
||||
self.__handle_system_error(err, module_id, module_name, method, **kwargs)
|
||||
return result
|
||||
|
||||
@@ -291,6 +294,7 @@ class ChainBase(metaclass=ABCMeta):
|
||||
# 中止继续执行
|
||||
break
|
||||
except Exception as err:
|
||||
logger.error(traceback.format_exc())
|
||||
self.__handle_system_error(err, module_id, module_name, method, **kwargs)
|
||||
return result
|
||||
|
||||
@@ -849,6 +853,8 @@ class ChainBase(metaclass=ABCMeta):
|
||||
:param kwargs: 其他参数(覆盖业务对象属性值)
|
||||
:return: 成功或失败
|
||||
"""
|
||||
# 添加格式化的时间参数
|
||||
kwargs.setdefault('current_time', datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
|
||||
# 渲染消息
|
||||
message = MessageTemplateHelper.render(message=message, meta=meta, mediainfo=mediainfo,
|
||||
torrentinfo=torrentinfo, transferinfo=transferinfo, **kwargs)
|
||||
@@ -932,6 +938,8 @@ class ChainBase(metaclass=ABCMeta):
|
||||
:param kwargs: 其他参数(覆盖业务对象属性值)
|
||||
:return: 成功或失败
|
||||
"""
|
||||
# 添加格式化的时间参数
|
||||
kwargs.setdefault('current_time', datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
|
||||
# 渲染消息
|
||||
message = MessageTemplateHelper.render(message=message, meta=meta, mediainfo=mediainfo,
|
||||
torrentinfo=torrentinfo, transferinfo=transferinfo, **kwargs)
|
||||
@@ -1055,6 +1063,18 @@ class ChainBase(metaclass=ABCMeta):
|
||||
"""
|
||||
return self.run_module("media_category")
|
||||
|
||||
def category_config(self) -> CategoryConfig:
|
||||
"""
|
||||
获取分类策略配置
|
||||
"""
|
||||
return self.run_module("load_category_config")
|
||||
|
||||
def save_category_config(self, config: CategoryConfig) -> bool:
|
||||
"""
|
||||
保存分类策略配置
|
||||
"""
|
||||
return self.run_module("save_category_config", config=config)
|
||||
|
||||
def register_commands(self, commands: Dict[str, dict]) -> None:
|
||||
"""
|
||||
注册菜单命令
|
||||
|
||||
318
app/chain/ai_recommend.py
Normal file
318
app/chain/ai_recommend.py
Normal file
@@ -0,0 +1,318 @@
|
||||
import re
|
||||
from typing import List, Optional, Dict, Any
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
|
||||
from app.chain import ChainBase
|
||||
from app.core.config import settings
|
||||
from app.log import logger
|
||||
from app.utils.common import log_execution_time
|
||||
from app.utils.singleton import Singleton
|
||||
from app.utils.string import StringUtils
|
||||
|
||||
|
||||
class AIRecommendChain(ChainBase, metaclass=Singleton):
|
||||
"""
|
||||
AI推荐处理链,单例运行
|
||||
用于基于搜索结果的AI智能推荐
|
||||
"""
|
||||
|
||||
# 缓存文件名
|
||||
__ai_indices_cache_file = "__ai_recommend_indices__"
|
||||
|
||||
# AI推荐状态
|
||||
_ai_recommend_running = False
|
||||
_ai_recommend_task: Optional[asyncio.Task] = None
|
||||
_current_request_hash: Optional[str] = None # 当前请求的哈希值
|
||||
_ai_recommend_result: Optional[List[int]] = None # AI推荐索引缓存(索引列表)
|
||||
_ai_recommend_error: Optional[str] = None # AI推荐错误信息
|
||||
|
||||
@staticmethod
|
||||
def _calculate_request_hash(
|
||||
filtered_indices: Optional[List[int]], search_results_count: int
|
||||
) -> str:
|
||||
"""
|
||||
计算请求的哈希值,用于判断请求是否变化
|
||||
"""
|
||||
request_data = {
|
||||
"filtered_indices": filtered_indices or [],
|
||||
"search_results_count": search_results_count,
|
||||
}
|
||||
return hashlib.md5(
|
||||
json.dumps(request_data, sort_keys=True).encode()
|
||||
).hexdigest()
|
||||
|
||||
@property
|
||||
def is_enabled(self) -> bool:
|
||||
"""
|
||||
检查AI推荐功能是否已启用。
|
||||
"""
|
||||
return settings.AI_AGENT_ENABLE and settings.AI_RECOMMEND_ENABLED
|
||||
|
||||
def _build_status(self) -> Dict[str, Any]:
|
||||
"""
|
||||
构建AI推荐状态字典
|
||||
:return: 状态字典
|
||||
"""
|
||||
if not self.is_enabled:
|
||||
return {"status": "disabled"}
|
||||
|
||||
if self._ai_recommend_running:
|
||||
return {"status": "running"}
|
||||
|
||||
# 尝试从数据库加载缓存
|
||||
if self._ai_recommend_result is None:
|
||||
cached_indices = self.load_cache(self.__ai_indices_cache_file)
|
||||
if cached_indices is not None:
|
||||
self._ai_recommend_result = cached_indices
|
||||
|
||||
# 只要有结果,始终返回completed状态和数据
|
||||
if self._ai_recommend_result is not None:
|
||||
return {"status": "completed", "results": self._ai_recommend_result}
|
||||
|
||||
if self._ai_recommend_error is not None:
|
||||
return {"status": "error", "error": self._ai_recommend_error}
|
||||
|
||||
return {"status": "idle"}
|
||||
|
||||
def get_current_status_only(self) -> Dict[str, Any]:
|
||||
"""
|
||||
获取当前状态(不校验hash,用于check_only模式)
|
||||
"""
|
||||
return self._build_status()
|
||||
|
||||
def get_status(
|
||||
self, filtered_indices: Optional[List[int]], search_results_count: int
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
获取AI推荐状态并检查请求是否变化(用于首次请求或force模式)
|
||||
如果请求变化(筛选条件变化),返回idle状态
|
||||
"""
|
||||
# 计算当前请求的hash
|
||||
request_hash = self._calculate_request_hash(
|
||||
filtered_indices, search_results_count
|
||||
)
|
||||
|
||||
# 检查请求是否变化
|
||||
is_same_request = request_hash == self._current_request_hash
|
||||
|
||||
# 如果请求变化了(筛选条件改变),返回idle状态
|
||||
if not is_same_request:
|
||||
return {"status": "idle"} if self.is_enabled else {"status": "disabled"}
|
||||
|
||||
# 请求未变化,返回当前实际状态
|
||||
return self._build_status()
|
||||
|
||||
@log_execution_time(logger=logger)
|
||||
async def async_ai_recommend(self, items: List[str], preference: str = None) -> str:
|
||||
"""
|
||||
AI推荐
|
||||
:param items: 候选资源列表(JSON字符串格式)
|
||||
:param preference: 用户偏好(可选)
|
||||
:return: AI返回的推荐结果
|
||||
"""
|
||||
# 设置运行状态
|
||||
self._ai_recommend_running = True
|
||||
try:
|
||||
# 导入LLMHelper
|
||||
from app.helper.llm import LLMHelper
|
||||
|
||||
# 获取LLM实例
|
||||
llm = LLMHelper.get_llm()
|
||||
|
||||
# 构建提示词
|
||||
user_preference = (
|
||||
preference
|
||||
or settings.AI_RECOMMEND_USER_PREFERENCE
|
||||
or "Prefer high-quality resources with more seeders"
|
||||
)
|
||||
|
||||
# 添加指令
|
||||
instruction = """
|
||||
Task: Select the best matching items from the list based on user preferences.
|
||||
|
||||
Each item contains:
|
||||
- index: Item number
|
||||
- title: Full torrent title
|
||||
- size: File size
|
||||
- seeders: Number of seeders
|
||||
|
||||
Output Format: Return ONLY a JSON array of "index" numbers (e.g., [0, 3, 1]). Do NOT include any explanations or other text.
|
||||
"""
|
||||
message = (
|
||||
f"User Preference: {user_preference}\n{instruction}\nCandidate Resources:\n"
|
||||
+ "\n".join(items)
|
||||
)
|
||||
|
||||
# 调用LLM
|
||||
response = await llm.ainvoke(message)
|
||||
return response.content
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"AI推荐配置错误: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
raise
|
||||
finally:
|
||||
# 清除运行状态
|
||||
self._ai_recommend_running = False
|
||||
self._ai_recommend_task = None
|
||||
|
||||
def is_ai_recommend_running(self) -> bool:
|
||||
"""
|
||||
检查AI推荐是否正在运行
|
||||
"""
|
||||
return self._ai_recommend_running
|
||||
|
||||
def cancel_ai_recommend(self):
|
||||
"""
|
||||
取消正在运行的AI推荐任务
|
||||
"""
|
||||
if self._ai_recommend_task and not self._ai_recommend_task.done():
|
||||
self._ai_recommend_task.cancel()
|
||||
self._ai_recommend_running = False
|
||||
self._ai_recommend_task = None
|
||||
self._current_request_hash = None
|
||||
self._ai_recommend_result = None
|
||||
self._ai_recommend_error = None
|
||||
self.remove_cache(self.__ai_indices_cache_file)
|
||||
|
||||
def start_recommend_task(
|
||||
self,
|
||||
filtered_indices: Optional[List[int]],
|
||||
search_results_count: int,
|
||||
results: List[Any],
|
||||
) -> None:
|
||||
"""
|
||||
启动AI推荐任务
|
||||
:param filtered_indices: 筛选后的索引列表
|
||||
:param search_results_count: 搜索结果总数
|
||||
:param results: 搜索结果列表
|
||||
"""
|
||||
# 防护检查:确保AI推荐功能已启用
|
||||
if not self.is_enabled:
|
||||
logger.warning("AI推荐功能未启用,跳过任务执行")
|
||||
return
|
||||
|
||||
# 计算新请求的哈希值
|
||||
new_request_hash = self._calculate_request_hash(
|
||||
filtered_indices, search_results_count
|
||||
)
|
||||
|
||||
# 如果请求变化了,取消旧任务
|
||||
if new_request_hash != self._current_request_hash:
|
||||
self.cancel_ai_recommend()
|
||||
|
||||
# 更新请求哈希值
|
||||
self._current_request_hash = new_request_hash
|
||||
|
||||
# 重置状态
|
||||
self._ai_recommend_result = None
|
||||
self._ai_recommend_error = None
|
||||
|
||||
# 启动新任务
|
||||
async def run_recommend():
|
||||
# 获取当前任务对象,用于在finally中比对
|
||||
current_task = asyncio.current_task()
|
||||
try:
|
||||
self._ai_recommend_running = True
|
||||
|
||||
# 准备数据
|
||||
items = []
|
||||
valid_indices = []
|
||||
max_items = settings.AI_RECOMMEND_MAX_ITEMS or 50
|
||||
|
||||
# 如果提供了筛选索引,先筛选结果;否则使用所有结果
|
||||
if filtered_indices is not None and len(filtered_indices) > 0:
|
||||
results_to_process = [
|
||||
results[i]
|
||||
for i in filtered_indices
|
||||
if 0 <= i < len(results)
|
||||
]
|
||||
else:
|
||||
results_to_process = results
|
||||
|
||||
for i, torrent in enumerate(results_to_process):
|
||||
if len(items) >= max_items:
|
||||
break
|
||||
|
||||
if not torrent.torrent_info:
|
||||
continue
|
||||
|
||||
valid_indices.append(i)
|
||||
|
||||
item_info = {
|
||||
"index": i,
|
||||
"title": torrent.torrent_info.title or "未知",
|
||||
"size": (
|
||||
StringUtils.format_size(torrent.torrent_info.size)
|
||||
if torrent.torrent_info.size
|
||||
else "0 B"
|
||||
),
|
||||
"seeders": torrent.torrent_info.seeders or 0,
|
||||
}
|
||||
|
||||
items.append(json.dumps(item_info, ensure_ascii=False))
|
||||
|
||||
if not items:
|
||||
self._ai_recommend_error = "没有可用于AI推荐的资源"
|
||||
return
|
||||
|
||||
# 调用AI推荐
|
||||
ai_response = await self.async_ai_recommend(items)
|
||||
|
||||
# 解析AI返回的索引
|
||||
try:
|
||||
# 使用正则提取JSON数组(非贪婪模式,避免匹配多个数组)
|
||||
json_match = re.search(r'\[.*?\]', ai_response, re.DOTALL)
|
||||
if not json_match:
|
||||
raise ValueError(ai_response)
|
||||
|
||||
ai_indices = json.loads(json_match.group())
|
||||
if not isinstance(ai_indices, list):
|
||||
raise ValueError(f"AI返回格式错误: {ai_response}")
|
||||
|
||||
# 映射回原始索引
|
||||
if filtered_indices:
|
||||
original_indices = [
|
||||
filtered_indices[valid_indices[i]]
|
||||
for i in ai_indices
|
||||
if i < len(valid_indices)
|
||||
and 0 <= filtered_indices[valid_indices[i]] < len(results)
|
||||
]
|
||||
else:
|
||||
original_indices = [
|
||||
valid_indices[i]
|
||||
for i in ai_indices
|
||||
if i < len(valid_indices)
|
||||
and 0 <= valid_indices[i] < len(results)
|
||||
]
|
||||
|
||||
# 只返回索引列表,不返回完整数据
|
||||
self._ai_recommend_result = original_indices
|
||||
|
||||
# 保存到数据库
|
||||
self.save_cache(original_indices, self.__ai_indices_cache_file)
|
||||
logger.info(f"AI推荐完成: {len(original_indices)}项")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"解析AI返回结果失败: {e}, 原始响应: {ai_response}"
|
||||
)
|
||||
self._ai_recommend_error = str(e)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("AI推荐任务被取消")
|
||||
except Exception as e:
|
||||
logger.error(f"AI推荐任务失败: {e}")
|
||||
self._ai_recommend_error = str(e)
|
||||
finally:
|
||||
# 只有当 self._ai_recommend_task 仍然是当前任务时,才清理状态
|
||||
# 如果任务被取消并启动了新任务,self._ai_recommend_task 已经指向新任务,不应重置
|
||||
if self._ai_recommend_task == current_task:
|
||||
self._ai_recommend_running = False
|
||||
self._ai_recommend_task = None
|
||||
|
||||
# 创建并启动任务
|
||||
self._ai_recommend_task = asyncio.create_task(run_recommend())
|
||||
@@ -152,7 +152,8 @@ class DownloadChain(ChainBase):
|
||||
save_path: Optional[str] = None,
|
||||
userid: Union[str, int] = None,
|
||||
username: Optional[str] = None,
|
||||
label: Optional[str] = None) -> Optional[str]:
|
||||
label: Optional[str] = None,
|
||||
return_detail: bool = False) -> Union[Optional[str], Tuple[Optional[str], Optional[str]]]:
|
||||
"""
|
||||
下载及发送通知
|
||||
:param context: 资源上下文
|
||||
@@ -166,6 +167,8 @@ class DownloadChain(ChainBase):
|
||||
:param userid: 用户ID
|
||||
:param username: 调用下载的用户名/插件名
|
||||
:param label: 自定义标签
|
||||
:param return_detail: 是否返回详细结果;False 时返回下载任务 hash 或 None,True 时返回 (hash, error_msg)
|
||||
:return: return_detail=False 时返回下载任务 hash 或 None;return_detail=True 时返回 (hash, error_msg)
|
||||
"""
|
||||
_torrent = context.torrent_info
|
||||
_media = context.media_info
|
||||
@@ -195,7 +198,7 @@ class DownloadChain(ChainBase):
|
||||
logger.debug(
|
||||
f"Resource download canceled by event: {event_data.source},"
|
||||
f"Reason: {event_data.reason}")
|
||||
return None
|
||||
return (None, "下载被事件取消") if return_detail else None
|
||||
# 如果事件修改了下载路径,使用新路径
|
||||
if event_data.options and event_data.options.get("save_path"):
|
||||
save_path = event_data.options.get("save_path")
|
||||
@@ -227,7 +230,7 @@ class DownloadChain(ChainBase):
|
||||
torrent_content = cache_backend.get(torrent_file.as_posix(), region="torrents")
|
||||
|
||||
if not torrent_content:
|
||||
return None
|
||||
return (None, "下载种子内容为空") if return_detail else None
|
||||
|
||||
# 获取种子文件的文件夹名和文件清单
|
||||
_folder_name, _file_list = TorrentHelper().get_fileinfo_from_torrent_content(torrent_content)
|
||||
@@ -259,7 +262,7 @@ class DownloadChain(ChainBase):
|
||||
logger.error(f"未找到下载目录:{_media.type.value} {_media.title_year}")
|
||||
self.messagehelper.put(f"{_media.type.value} {_media.title_year} 未找到下载目录!",
|
||||
title="下载失败", role="system")
|
||||
return None
|
||||
return (None, "未找到下载目录") if return_detail else None
|
||||
fileURI = FileURI(storage=storage, path=download_dir.as_posix())
|
||||
download_dir = Path(fileURI.uri)
|
||||
|
||||
@@ -327,9 +330,10 @@ class DownloadChain(ChainBase):
|
||||
if not file_meta.begin_episode \
|
||||
or file_meta.begin_episode not in episodes:
|
||||
continue
|
||||
# 只处理视频格式
|
||||
# 只处理音视频、字幕格式
|
||||
media_exts = settings.RMT_MEDIAEXT + settings.RMT_SUBEXT + settings.RMT_AUDIOEXT
|
||||
if not Path(file).suffix \
|
||||
or Path(file).suffix.lower() not in settings.RMT_MEDIAEXT:
|
||||
or Path(file).suffix.lower() not in media_exts:
|
||||
continue
|
||||
files_to_add.append({
|
||||
"download_hash": _hash,
|
||||
@@ -387,6 +391,8 @@ class DownloadChain(ChainBase):
|
||||
f"错误信息:{error_msg}",
|
||||
image=_media.get_message_image(),
|
||||
userid=userid))
|
||||
if return_detail:
|
||||
return _hash, error_msg
|
||||
return _hash
|
||||
|
||||
def batch_download(self,
|
||||
|
||||
1223
app/chain/media.py
1223
app/chain/media.py
File diff suppressed because it is too large
Load Diff
@@ -40,7 +40,7 @@ class MessageChain(ChainBase):
|
||||
# 用户会话信息 {userid: (session_id, last_time)}
|
||||
_user_sessions: Dict[Union[str, int], tuple] = {}
|
||||
# 会话超时时间(分钟)
|
||||
_session_timeout_minutes: int = 15
|
||||
_session_timeout_minutes: int = 24 * 60
|
||||
|
||||
@staticmethod
|
||||
def __get_noexits_info(
|
||||
@@ -112,8 +112,8 @@ class MessageChain(ChainBase):
|
||||
channel = info.channel
|
||||
# 用户ID
|
||||
userid = info.userid
|
||||
# 用户名
|
||||
username = info.username or userid
|
||||
# 用户名(当渠道未提供公开用户名时,回退为 userid 的字符串,避免后续类型校验异常)
|
||||
username = str(info.username) if info.username not in (None, "") else str(userid)
|
||||
if userid is None or userid == '':
|
||||
logger.debug(f'未识别到用户ID:{body}{form}{args}')
|
||||
return
|
||||
@@ -195,10 +195,14 @@ class MessageChain(ChainBase):
|
||||
if text.isdigit():
|
||||
# 用户选择了具体的条目
|
||||
# 缓存
|
||||
cache_data: dict = user_cache.get(userid).copy()
|
||||
cache_data: dict = user_cache.get(userid)
|
||||
if not cache_data:
|
||||
# 发送消息
|
||||
self.post_message(Notification(channel=channel, source=source, title="输入有误!", userid=userid))
|
||||
return
|
||||
cache_data = cache_data.copy()
|
||||
# 选择项目
|
||||
if not cache_data \
|
||||
or not cache_data.get('items') \
|
||||
if not cache_data.get('items') \
|
||||
or len(cache_data.get('items')) < int(text):
|
||||
# 发送消息
|
||||
self.post_message(Notification(channel=channel, source=source, title="输入有误!", userid=userid))
|
||||
@@ -370,12 +374,13 @@ class MessageChain(ChainBase):
|
||||
del cache_data
|
||||
elif text.lower() == "p":
|
||||
# 上一页
|
||||
cache_data: dict = user_cache.get(userid).copy()
|
||||
cache_data: dict = user_cache.get(userid)
|
||||
if not cache_data:
|
||||
# 没有缓存
|
||||
self.post_message(Notification(
|
||||
channel=channel, source=source, title="输入有误!", userid=userid))
|
||||
return
|
||||
cache_data = cache_data.copy()
|
||||
try:
|
||||
if _current_page == 0:
|
||||
# 第一页
|
||||
@@ -422,12 +427,13 @@ class MessageChain(ChainBase):
|
||||
del cache_data
|
||||
elif text.lower() == "n":
|
||||
# 下一页
|
||||
cache_data: dict = user_cache.get(userid).copy()
|
||||
cache_data: dict = user_cache.get(userid)
|
||||
if not cache_data:
|
||||
# 没有缓存
|
||||
self.post_message(Notification(
|
||||
channel=channel, source=source, title="输入有误!", userid=userid))
|
||||
return
|
||||
cache_data = cache_data.copy()
|
||||
try:
|
||||
cache_type: str = cache_data.get('type')
|
||||
# 产生副本,避免修改原值
|
||||
@@ -484,18 +490,14 @@ class MessageChain(ChainBase):
|
||||
# 重新搜索/下载
|
||||
content = re.sub(r"(搜索|下载)[::\s]*", "", text)
|
||||
action = "ReSearch"
|
||||
elif text.startswith("#") \
|
||||
or re.search(r"^请[问帮你]", text) \
|
||||
or re.search(r"[??]$", text) \
|
||||
or StringUtils.count_words(text) > 10 \
|
||||
or text.find("继续") != -1:
|
||||
# 聊天
|
||||
content = text
|
||||
action = "Chat"
|
||||
elif StringUtils.is_link(text):
|
||||
# 链接
|
||||
content = text
|
||||
action = "Link"
|
||||
elif not StringUtils.is_media_title_like(text):
|
||||
# 聊天
|
||||
content = text
|
||||
action = "Chat"
|
||||
else:
|
||||
# 搜索
|
||||
content = text
|
||||
@@ -836,8 +838,7 @@ class MessageChain(ChainBase):
|
||||
|
||||
return buttons
|
||||
|
||||
@staticmethod
|
||||
def _get_or_create_session_id(userid: Union[str, int]) -> str:
|
||||
def _get_or_create_session_id(self, userid: Union[str, int]) -> str:
|
||||
"""
|
||||
获取或创建会话ID
|
||||
如果用户上次会话在15分钟内,则复用相同的会话ID;否则创建新的会话ID
|
||||
@@ -845,34 +846,33 @@ class MessageChain(ChainBase):
|
||||
current_time = datetime.now()
|
||||
|
||||
# 检查用户是否有已存在的会话
|
||||
if userid in MessageChain._user_sessions:
|
||||
session_id, last_time = MessageChain._user_sessions[userid]
|
||||
if userid in self._user_sessions:
|
||||
session_id, last_time = self._user_sessions[userid]
|
||||
|
||||
# 计算时间差
|
||||
time_diff = current_time - last_time
|
||||
|
||||
# 如果时间差小于等于15分钟,复用会话ID
|
||||
if time_diff <= timedelta(minutes=MessageChain._session_timeout_minutes):
|
||||
# 如果时间差小于等于xx分钟,复用会话ID
|
||||
if time_diff <= timedelta(minutes=self._session_timeout_minutes):
|
||||
# 更新最后使用时间
|
||||
MessageChain._user_sessions[userid] = (session_id, current_time)
|
||||
self._user_sessions[userid] = (session_id, current_time)
|
||||
logger.info(
|
||||
f"复用会话ID: {session_id}, 用户: {userid}, 距离上次会话: {time_diff.total_seconds() / 60:.1f}分钟")
|
||||
return session_id
|
||||
|
||||
# 创建新的会话ID
|
||||
new_session_id = f"user_{userid}_{int(time.time())}"
|
||||
MessageChain._user_sessions[userid] = (new_session_id, current_time)
|
||||
self._user_sessions[userid] = (new_session_id, current_time)
|
||||
logger.info(f"创建新会话ID: {new_session_id}, 用户: {userid}")
|
||||
return new_session_id
|
||||
|
||||
@staticmethod
|
||||
def clear_user_session(userid: Union[str, int]) -> bool:
|
||||
def clear_user_session(self, userid: Union[str, int]) -> bool:
|
||||
"""
|
||||
清除指定用户的会话信息
|
||||
返回是否成功清除
|
||||
"""
|
||||
if userid in MessageChain._user_sessions:
|
||||
session_id, _ = MessageChain._user_sessions.pop(userid)
|
||||
if userid in self._user_sessions:
|
||||
session_id, _ = self._user_sessions.pop(userid)
|
||||
logger.info(f"已清除用户 {userid} 的会话: {session_id}")
|
||||
return True
|
||||
return False
|
||||
@@ -883,8 +883,8 @@ class MessageChain(ChainBase):
|
||||
"""
|
||||
# 获取并清除会话信息
|
||||
session_id = None
|
||||
if userid in MessageChain._user_sessions:
|
||||
session_id, _ = MessageChain._user_sessions.pop(userid)
|
||||
if userid in self._user_sessions:
|
||||
session_id, _ = self._user_sessions.pop(userid)
|
||||
logger.info(f"已清除用户 {userid} 的会话: {session_id}")
|
||||
|
||||
# 如果有会话ID,同时清除智能体的会话记忆
|
||||
|
||||
@@ -6,7 +6,7 @@ from app.chain import ChainBase
|
||||
from app.chain.bangumi import BangumiChain
|
||||
from app.chain.douban import DoubanChain
|
||||
from app.chain.tmdb import TmdbChain
|
||||
from app.core.cache import cached
|
||||
from app.core.cache import cached, fresh
|
||||
from app.core.config import settings, global_vars
|
||||
from app.helper.image import ImageHelper
|
||||
from app.log import logger
|
||||
@@ -27,9 +27,11 @@ class RecommendChain(ChainBase, metaclass=Singleton):
|
||||
# 推荐缓存区域
|
||||
recommend_cache_region = "recommend"
|
||||
|
||||
def refresh_recommend(self):
|
||||
def refresh_recommend(self, manual: bool = False):
|
||||
"""
|
||||
刷新推荐
|
||||
|
||||
:param manual: 手动触发
|
||||
"""
|
||||
logger.debug("Starting to refresh Recommend data.")
|
||||
|
||||
@@ -62,7 +64,9 @@ class RecommendChain(ChainBase, metaclass=Singleton):
|
||||
if method in methods_finished:
|
||||
continue
|
||||
logger.debug(f"Fetch {method.__name__} data for page {page}.")
|
||||
data = method(page=page)
|
||||
# 手动触发的刷新,总是需要获取最新数据
|
||||
with fresh(manual):
|
||||
data = method(page=page)
|
||||
if not data:
|
||||
logger.debug("All recommendation methods have finished fetching data. Ending pagination early.")
|
||||
methods_finished.add(method)
|
||||
@@ -90,7 +94,6 @@ class RecommendChain(ChainBase, metaclass=Singleton):
|
||||
poster_path = data.get("poster_path")
|
||||
if poster_path:
|
||||
poster_url = poster_path.replace("original", "w500")
|
||||
logger.debug(f"Caching poster image: {poster_url}")
|
||||
self.__fetch_and_save_image(poster_url)
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -29,6 +29,7 @@ class SearchChain(ChainBase):
|
||||
"""
|
||||
|
||||
__result_temp_file = "__search_result__"
|
||||
__ai_result_temp_file = "__ai_search_result__"
|
||||
|
||||
def search_by_id(self, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||||
mtype: MediaType = None, area: Optional[str] = "title", season: Optional[int] = None,
|
||||
@@ -48,7 +49,7 @@ class SearchChain(ChainBase):
|
||||
logger.error(f'{tmdbid} 媒体信息识别失败!')
|
||||
return []
|
||||
no_exists = None
|
||||
if season:
|
||||
if season is not None:
|
||||
no_exists = {
|
||||
tmdbid or doubanid: {
|
||||
season: NotExistMediaInfo(episodes=[])
|
||||
@@ -98,6 +99,18 @@ class SearchChain(ChainBase):
|
||||
"""
|
||||
return await self.async_load_cache(self.__result_temp_file)
|
||||
|
||||
async def async_last_ai_results(self) -> Optional[List[Context]]:
|
||||
"""
|
||||
异步获取上次AI推荐结果
|
||||
"""
|
||||
return await self.async_load_cache(self.__ai_result_temp_file)
|
||||
|
||||
async def async_save_ai_results(self, results: List[Context]):
|
||||
"""
|
||||
异步保存AI推荐结果
|
||||
"""
|
||||
await self.async_save_cache(results, self.__ai_result_temp_file)
|
||||
|
||||
async def async_search_by_id(self, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||||
mtype: MediaType = None, area: Optional[str] = "title", season: Optional[int] = None,
|
||||
sites: List[int] = None, cache_local: bool = False) -> List[Context]:
|
||||
@@ -116,7 +129,7 @@ class SearchChain(ChainBase):
|
||||
logger.error(f'{tmdbid} 媒体信息识别失败!')
|
||||
return []
|
||||
no_exists = None
|
||||
if season:
|
||||
if season is not None:
|
||||
no_exists = {
|
||||
tmdbid or doubanid: {
|
||||
season: NotExistMediaInfo(episodes=[])
|
||||
@@ -168,7 +181,7 @@ class SearchChain(ChainBase):
|
||||
# 过滤剧集
|
||||
season_episodes = {sea: info.episodes
|
||||
for sea, info in no_exists[mediakey].items()}
|
||||
elif mediainfo.season:
|
||||
elif mediainfo.season is not None:
|
||||
# 豆瓣只搜索当前季
|
||||
season_episodes = {mediainfo.season: []}
|
||||
else:
|
||||
@@ -267,7 +280,7 @@ class SearchChain(ChainBase):
|
||||
logger.info(f"种子名称应用识别词后发生改变:{torrent.title} => {torrent_meta.org_string}")
|
||||
# 季集数过滤
|
||||
if season_episodes \
|
||||
and not torrenthelper.match_season_episodes(torrent=torrent,
|
||||
and not TorrentHelper.match_season_episodes(torrent=torrent,
|
||||
meta=torrent_meta,
|
||||
season_episodes=season_episodes):
|
||||
continue
|
||||
|
||||
@@ -44,6 +44,7 @@ class SiteChain(ChainBase):
|
||||
"star-space.net": self.__indexphp_test,
|
||||
"yemapt.org": self.__yema_test,
|
||||
"hddolby.com": self.__hddolby_test,
|
||||
"rousi.pro": self.__rousi_test,
|
||||
}
|
||||
|
||||
def refresh_userdata(self, site: dict = None) -> Optional[SiteUserData]:
|
||||
@@ -249,6 +250,32 @@ class SiteChain(ChainBase):
|
||||
else:
|
||||
return False, f"错误:{res.status_code} {res.reason}"
|
||||
|
||||
@staticmethod
|
||||
def __rousi_test(site: Site) -> Tuple[bool, str]:
|
||||
"""
|
||||
判断站点是否已经登陆:rousi
|
||||
"""
|
||||
url = f"https://{StringUtils.get_url_domain(site.url)}/api/v1/profile"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
"Authorization": f"Bearer {site.apikey}",
|
||||
}
|
||||
res = RequestUtils(
|
||||
headers=headers,
|
||||
proxies=settings.PROXY if site.proxy else None,
|
||||
timeout=site.timeout or 15
|
||||
).get_res(url=url)
|
||||
if res is None:
|
||||
return False, "无法打开网站!"
|
||||
if res.status_code == 200:
|
||||
user_info = res.json()
|
||||
if user_info and user_info.get("code") == 0:
|
||||
return True, "连接成功"
|
||||
return False, "APIKEY已过期"
|
||||
else:
|
||||
return False, f"错误:{res.status_code} {res.reason}"
|
||||
|
||||
@staticmethod
|
||||
def __parse_favicon(url: str, cookie: str, ua: str) -> Tuple[str, Optional[str]]:
|
||||
"""
|
||||
@@ -462,20 +489,18 @@ class SiteChain(ChainBase):
|
||||
logger.warn(f"站点 {domain} 索引器不存在!")
|
||||
return
|
||||
# 查询站点图标
|
||||
site_icon = siteoper.get_icon_by_domain(domain)
|
||||
if not site_icon or not site_icon.base64:
|
||||
logger.info(f"开始缓存站点 {indexer.get('name')} 图标 ...")
|
||||
icon_url, icon_base64 = self.__parse_favicon(url=indexer.get("domain"),
|
||||
cookie=cookie,
|
||||
ua=settings.USER_AGENT)
|
||||
if icon_url:
|
||||
siteoper.update_icon(name=indexer.get("name"),
|
||||
domain=domain,
|
||||
icon_url=icon_url,
|
||||
icon_base64=icon_base64)
|
||||
logger.info(f"缓存站点 {indexer.get('name')} 图标成功")
|
||||
else:
|
||||
logger.warn(f"缓存站点 {indexer.get('name')} 图标失败")
|
||||
logger.info(f"开始缓存站点 {indexer.get('name')} 图标 ...")
|
||||
icon_url, icon_base64 = self.__parse_favicon(url=indexer.get("domain"),
|
||||
cookie=cookie,
|
||||
ua=settings.USER_AGENT)
|
||||
if icon_url:
|
||||
siteoper.update_icon(name=indexer.get("name"),
|
||||
domain=domain,
|
||||
icon_url=icon_url,
|
||||
icon_base64=icon_base64)
|
||||
logger.info(f"缓存站点 {indexer.get('name')} 图标成功")
|
||||
else:
|
||||
logger.warn(f"缓存站点 {indexer.get('name')} 图标失败")
|
||||
|
||||
@eventmanager.register(EventType.SiteUpdated)
|
||||
def clear_site_data(self, event: Event):
|
||||
|
||||
@@ -31,6 +31,12 @@ class StorageChain(ChainBase):
|
||||
"""
|
||||
return self.run_module("generate_qrcode", storage=storage)
|
||||
|
||||
def generate_auth_url(self, storage: str) -> Optional[Tuple[dict, str]]:
|
||||
"""
|
||||
生成 OAuth2 授权 URL
|
||||
"""
|
||||
return self.run_module("generate_auth_url", storage=storage)
|
||||
|
||||
def check_login(self, storage: str, **kwargs) -> Optional[Tuple[dict, str]]:
|
||||
"""
|
||||
登录确认
|
||||
@@ -133,30 +139,41 @@ class StorageChain(ChainBase):
|
||||
"""
|
||||
return self.run_module("support_transtype", storage=storage)
|
||||
|
||||
def is_bluray_folder(self, fileitem: Optional[schemas.FileItem]) -> bool:
|
||||
"""
|
||||
检查是否蓝光目录
|
||||
"""
|
||||
if not fileitem or fileitem.type != "dir":
|
||||
return False
|
||||
if self.get_file_item(storage=fileitem.storage, path=Path(fileitem.path) / "BDMV"):
|
||||
return True
|
||||
if self.get_file_item(storage=fileitem.storage, path=Path(fileitem.path) / "CERTIFICATE"):
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def contains_bluray_subdirectories(fileitems: Optional[List[schemas.FileItem]]) -> bool:
|
||||
"""
|
||||
判断是否包含蓝光必备的文件夹
|
||||
"""
|
||||
required_files = {"BDMV", "CERTIFICATE"}
|
||||
return any(
|
||||
item.type == "dir" and item.name in required_files
|
||||
for item in fileitems or []
|
||||
)
|
||||
|
||||
def delete_media_file(self, fileitem: schemas.FileItem, delete_self: bool = True) -> bool:
|
||||
"""
|
||||
删除媒体文件,以及不含媒体文件的目录
|
||||
"""
|
||||
|
||||
def __is_bluray_dir(_fileitem: schemas.FileItem) -> bool:
|
||||
"""
|
||||
检查是否蓝光目录
|
||||
"""
|
||||
_dir_files = self.list_files(fileitem=_fileitem, recursion=False)
|
||||
if _dir_files:
|
||||
for _f in _dir_files:
|
||||
if _f.type == "dir" and _f.name in ["BDMV", "CERTIFICATE"]:
|
||||
return True
|
||||
return False
|
||||
|
||||
media_exts = settings.RMT_MEDIAEXT + settings.DOWNLOAD_TMPEXT
|
||||
media_exts = settings.RMT_MEDIAEXT + settings.DOWNLOAD_TMPEXT + settings.RMT_SUBEXT + settings.RMT_AUDIOEXT
|
||||
fileitem_path = Path(fileitem.path) if fileitem.path else Path("")
|
||||
if len(fileitem_path.parts) <= 2:
|
||||
logger.warn(f"【{fileitem.storage}】{fileitem.path} 根目录或一级目录不允许删除")
|
||||
return False
|
||||
if fileitem.type == "dir":
|
||||
# 本身是目录
|
||||
if __is_bluray_dir(fileitem):
|
||||
if self.is_bluray_folder(fileitem):
|
||||
logger.warn(f"正在删除蓝光原盘目录:【{fileitem.storage}】{fileitem.path}")
|
||||
if not self.delete_file(fileitem):
|
||||
logger.warn(f"【{fileitem.storage}】{fileitem.path} 删除失败")
|
||||
|
||||
@@ -144,7 +144,7 @@ class SubscribeChain(ChainBase):
|
||||
metainfo.year = year
|
||||
if mtype:
|
||||
metainfo.type = mtype
|
||||
if season:
|
||||
if season is not None:
|
||||
metainfo.type = MediaType.TV
|
||||
metainfo.begin_season = season
|
||||
# 识别媒体信息
|
||||
@@ -174,7 +174,7 @@ class SubscribeChain(ChainBase):
|
||||
# 豆瓣标题处理
|
||||
meta = MetaInfo(mediainfo.title)
|
||||
mediainfo.title = meta.name
|
||||
if not season:
|
||||
if season is None:
|
||||
season = meta.begin_season
|
||||
|
||||
# 使用名称识别兜底
|
||||
@@ -188,7 +188,7 @@ class SubscribeChain(ChainBase):
|
||||
|
||||
# 总集数
|
||||
if mediainfo.type == MediaType.TV:
|
||||
if not season:
|
||||
if season is None:
|
||||
season = 1
|
||||
# 总集数
|
||||
if not kwargs.get('total_episode'):
|
||||
@@ -292,7 +292,7 @@ class SubscribeChain(ChainBase):
|
||||
"description": mediainfo.overview
|
||||
})
|
||||
# 返回结果
|
||||
return sid, ""
|
||||
return sid, err_msg
|
||||
|
||||
async def async_add(self, title: str, year: str,
|
||||
mtype: MediaType = None,
|
||||
@@ -321,7 +321,7 @@ class SubscribeChain(ChainBase):
|
||||
metainfo.year = year
|
||||
if mtype:
|
||||
metainfo.type = mtype
|
||||
if season:
|
||||
if season is not None:
|
||||
metainfo.type = MediaType.TV
|
||||
metainfo.begin_season = season
|
||||
# 识别媒体信息
|
||||
@@ -351,7 +351,7 @@ class SubscribeChain(ChainBase):
|
||||
# 豆瓣标题处理
|
||||
meta = MetaInfo(mediainfo.title)
|
||||
mediainfo.title = meta.name
|
||||
if not season:
|
||||
if season is None:
|
||||
season = meta.begin_season
|
||||
|
||||
# 使用名称识别兜底
|
||||
@@ -365,7 +365,7 @@ class SubscribeChain(ChainBase):
|
||||
|
||||
# 总集数
|
||||
if mediainfo.type == MediaType.TV:
|
||||
if not season:
|
||||
if season is None:
|
||||
season = 1
|
||||
# 总集数
|
||||
if not kwargs.get('total_episode'):
|
||||
@@ -469,7 +469,7 @@ class SubscribeChain(ChainBase):
|
||||
"description": mediainfo.overview
|
||||
})
|
||||
# 返回结果
|
||||
return sid, ""
|
||||
return sid, err_msg
|
||||
|
||||
@staticmethod
|
||||
def exists(mediainfo: MediaInfo, meta: MetaBase = None):
|
||||
@@ -530,7 +530,7 @@ class SubscribeChain(ChainBase):
|
||||
# 生成元数据
|
||||
meta = MetaInfo(subscribe.name)
|
||||
meta.year = subscribe.year
|
||||
meta.begin_season = subscribe.season or None
|
||||
meta.begin_season = subscribe.season if subscribe.season is not None else None
|
||||
try:
|
||||
meta.type = MediaType(subscribe.type)
|
||||
except ValueError:
|
||||
@@ -1119,6 +1119,19 @@ class SubscribeChain(ChainBase):
|
||||
})
|
||||
logger.info(f'{subscribe.name} 订阅元数据更新完成')
|
||||
|
||||
def get_subscribe_by_source(self, source: str) -> Optional[Subscribe]:
|
||||
"""
|
||||
从来源获取订阅
|
||||
"""
|
||||
source_keyword = self.parse_subscribe_source_keyword(source)
|
||||
if not source_keyword:
|
||||
return None
|
||||
# 只保留需要的字段动态获取订阅
|
||||
valid_fields = {k: v for k, v in source_keyword.items()
|
||||
if k in ["type", "season", "tmdbid", "doubanid", "bangumiid"]}
|
||||
# 暂时不考虑订阅历史, 若有必要再添加
|
||||
return SubscribeOper().get_by(**valid_fields)
|
||||
|
||||
@staticmethod
|
||||
def follow():
|
||||
"""
|
||||
@@ -1635,7 +1648,7 @@ class SubscribeChain(ChainBase):
|
||||
info = schemas.SubscribeEpisodeInfo()
|
||||
info.title = episode.name
|
||||
info.description = episode.overview
|
||||
info.backdrop = f"https://{settings.TMDB_IMAGE_DOMAIN}/t/p/w500${episode.still_path}"
|
||||
info.backdrop = settings.TMDB_IMAGE_URL(episode.still_path, "w500")
|
||||
episodes[episode.episode_number] = info
|
||||
elif subscribe.type == MediaType.TV.value:
|
||||
# 根据开始结束集计算集信息
|
||||
@@ -1655,7 +1668,7 @@ class SubscribeChain(ChainBase):
|
||||
if download_his:
|
||||
for his in download_his:
|
||||
# 查询下载文件
|
||||
files = downloadhis.get_files_by_hash(his.download_hash)
|
||||
files = downloadhis.get_files_by_hash(his.download_hash, state=1)
|
||||
if files:
|
||||
for file in files:
|
||||
# 识别文件名
|
||||
@@ -1828,8 +1841,9 @@ class SubscribeChain(ChainBase):
|
||||
def get_subscribe_source_keyword(subscribe: Subscribe) -> str:
|
||||
"""
|
||||
构造用于订阅来源的关键字字符串
|
||||
|
||||
:param subscribe: Subscribe 对象
|
||||
:return: 格式化的订阅来源关键字字符串,格式为 "Subscribe|{...}"
|
||||
:return str: 格式化的订阅来源关键字字符串,格式为 "Subscribe|{...}"
|
||||
"""
|
||||
source_keyword = {
|
||||
'id': subscribe.id,
|
||||
@@ -1844,3 +1858,24 @@ class SubscribeChain(ChainBase):
|
||||
'bangumiid': subscribe.bangumiid
|
||||
}
|
||||
return f"Subscribe|{json.dumps(source_keyword, ensure_ascii=False)}"
|
||||
|
||||
@staticmethod
|
||||
def parse_subscribe_source_keyword(source_keyword_str: str) -> Optional[dict]:
|
||||
"""
|
||||
解析订阅来源关键字字符串
|
||||
|
||||
:param source_keyword_str: 订阅来源关键字字符串,格式为 "Subscribe|{...}"
|
||||
:return Dict: 如果解析失败则返回None
|
||||
"""
|
||||
if not source_keyword_str or not source_keyword_str.startswith("Subscribe|"):
|
||||
return None
|
||||
|
||||
try:
|
||||
# 分割字符串获取JSON部分
|
||||
json_part = source_keyword_str.split("|", 1)[1]
|
||||
# 解析JSON字符串
|
||||
source_keyword = json.loads(json_part)
|
||||
return source_keyword
|
||||
except (IndexError, json.JSONDecodeError, TypeError) as e:
|
||||
logger.error(f"解析订阅来源关键字失败: {e}")
|
||||
return None
|
||||
|
||||
@@ -265,6 +265,9 @@ class TorrentsChain(ChainBase):
|
||||
for torrent in torrents:
|
||||
if global_vars.is_system_stopped:
|
||||
break
|
||||
if not torrent.enclosure:
|
||||
logger.warn(f"缺少种子链接,忽略处理: {torrent.title}")
|
||||
continue
|
||||
logger.info(f'处理资源:{torrent.title} ...')
|
||||
# 识别
|
||||
meta = MetaInfo(title=torrent.title, subtitle=torrent.description)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -27,8 +27,6 @@ DEFAULT_CACHE_SIZE = 1024
|
||||
# 默认缓存有效期
|
||||
DEFAULT_CACHE_TTL = 365 * 24 * 60 * 60
|
||||
|
||||
lock = threading.Lock()
|
||||
|
||||
# 上下文变量来控制缓存行为
|
||||
_fresh = contextvars.ContextVar('fresh', default=False)
|
||||
|
||||
@@ -297,14 +295,14 @@ class AsyncCacheBackend(CacheBackend):
|
||||
"""
|
||||
获取所有缓存键,类似 dict.keys()(异步)
|
||||
"""
|
||||
async for key, _ in await self.items(region=region):
|
||||
async for key, _ in self.items(region=region):
|
||||
yield key
|
||||
|
||||
async def values(self, region: Optional[str] = DEFAULT_CACHE_REGION) -> AsyncGenerator[Any, None]:
|
||||
"""
|
||||
获取所有缓存值,类似 dict.values()(异步)
|
||||
"""
|
||||
async for _, value in await self.items(region=region):
|
||||
async for _, value in self.items(region=region):
|
||||
yield value
|
||||
|
||||
async def update(self, other: Dict[str, Any], region: Optional[str] = DEFAULT_CACHE_REGION,
|
||||
@@ -332,7 +330,7 @@ class AsyncCacheBackend(CacheBackend):
|
||||
弹出最后一个缓存项,类似 dict.popitem()(异步)
|
||||
"""
|
||||
items = []
|
||||
async for item in await self.items(region=region):
|
||||
async for item in self.items(region=region):
|
||||
items.append(item)
|
||||
if not items:
|
||||
raise KeyError("popitem(): cache is empty")
|
||||
@@ -364,6 +362,11 @@ class MemoryBackend(CacheBackend):
|
||||
基于 `cachetools.TTLCache` 实现的缓存后端
|
||||
"""
|
||||
|
||||
# 类变量 _region_caches 的互斥锁
|
||||
_lock = threading.Lock()
|
||||
# 存储各个 region 的缓存实例,region -> TTLCache
|
||||
_region_caches: Dict[str, Union[MemoryTTLCache, MemoryLRUCache]] = {}
|
||||
|
||||
def __init__(self, cache_type: Literal['ttl', 'lru'] = 'ttl',
|
||||
maxsize: Optional[int] = None, ttl: Optional[int] = None):
|
||||
"""
|
||||
@@ -376,8 +379,6 @@ class MemoryBackend(CacheBackend):
|
||||
self.cache_type = cache_type
|
||||
self.maxsize = maxsize or DEFAULT_CACHE_SIZE
|
||||
self.ttl = ttl or DEFAULT_CACHE_TTL
|
||||
# 存储各个 region 的缓存实例,region -> TTLCache
|
||||
self._region_caches: Dict[str, Union[MemoryTTLCache, MemoryLRUCache]] = {}
|
||||
|
||||
def __get_region_cache(self, region: str) -> Optional[Union[MemoryTTLCache, MemoryLRUCache]]:
|
||||
"""
|
||||
@@ -400,7 +401,7 @@ class MemoryBackend(CacheBackend):
|
||||
maxsize = kwargs.get("maxsize", self.maxsize)
|
||||
region = self.get_region(region)
|
||||
# 设置缓存值
|
||||
with lock:
|
||||
with self._lock:
|
||||
# 如果该 key 尚未有缓存实例,则创建一个新的 TTLCache 实例
|
||||
region_cache = self._region_caches.setdefault(
|
||||
region,
|
||||
@@ -445,7 +446,7 @@ class MemoryBackend(CacheBackend):
|
||||
region_cache = self.__get_region_cache(region)
|
||||
if region_cache is None:
|
||||
return
|
||||
with lock:
|
||||
with self._lock:
|
||||
del region_cache[key]
|
||||
|
||||
def clear(self, region: Optional[str] = DEFAULT_CACHE_REGION) -> None:
|
||||
@@ -458,13 +459,13 @@ class MemoryBackend(CacheBackend):
|
||||
# 清理指定缓存区
|
||||
region_cache = self.__get_region_cache(region)
|
||||
if region_cache:
|
||||
with lock:
|
||||
with self._lock:
|
||||
region_cache.clear()
|
||||
logger.debug(f"Cleared cache for region: {region}")
|
||||
else:
|
||||
# 清除所有区域的缓存
|
||||
for region_cache in self._region_caches.values():
|
||||
with lock:
|
||||
with self._lock:
|
||||
region_cache.clear()
|
||||
logger.info("Cleared all cache")
|
||||
|
||||
@@ -480,7 +481,7 @@ class MemoryBackend(CacheBackend):
|
||||
yield from ()
|
||||
return
|
||||
# 使用锁保护迭代过程,避免在迭代时缓存被修改
|
||||
with lock:
|
||||
with self._lock:
|
||||
# 创建快照避免并发修改问题
|
||||
items_snapshot = list(region_cache.items())
|
||||
for item in items_snapshot:
|
||||
@@ -507,18 +508,7 @@ class AsyncMemoryBackend(AsyncCacheBackend):
|
||||
:param maxsize: 缓存的最大条目数
|
||||
:param ttl: 默认缓存存活时间,单位秒
|
||||
"""
|
||||
self.cache_type = cache_type
|
||||
self.maxsize = maxsize or DEFAULT_CACHE_SIZE
|
||||
self.ttl = ttl or DEFAULT_CACHE_TTL
|
||||
# 存储各个 region 的缓存实例,region -> TTLCache
|
||||
self._region_caches: Dict[str, Union[MemoryTTLCache, MemoryLRUCache]] = {}
|
||||
|
||||
def __get_region_cache(self, region: str) -> Optional[Union[MemoryTTLCache, MemoryLRUCache]]:
|
||||
"""
|
||||
获取指定区域的缓存实例,如果不存在则返回 None
|
||||
"""
|
||||
region = self.get_region(region)
|
||||
return self._region_caches.get(region)
|
||||
self._backend = MemoryBackend(cache_type=cache_type, maxsize=maxsize, ttl=ttl)
|
||||
|
||||
async def set(self, key: str, value: Any, ttl: Optional[int] = None,
|
||||
region: Optional[str] = DEFAULT_CACHE_REGION, **kwargs) -> None:
|
||||
@@ -530,18 +520,7 @@ class AsyncMemoryBackend(AsyncCacheBackend):
|
||||
:param ttl: 缓存的存活时间,不传入为永久缓存,单位秒
|
||||
:param region: 缓存的区
|
||||
"""
|
||||
ttl = ttl or self.ttl
|
||||
maxsize = kwargs.get("maxsize", self.maxsize)
|
||||
region = self.get_region(region)
|
||||
# 设置缓存值
|
||||
with lock:
|
||||
# 如果该 key 尚未有缓存实例,则创建一个新的 TTLCache 实例
|
||||
region_cache = self._region_caches.setdefault(
|
||||
region,
|
||||
MemoryTTLCache(maxsize=maxsize, ttl=ttl) if self.cache_type == 'ttl'
|
||||
else MemoryLRUCache(maxsize=maxsize)
|
||||
)
|
||||
region_cache[key] = value
|
||||
return self._backend.set(key=key, value=value, ttl=ttl, region=region, **kwargs)
|
||||
|
||||
async def exists(self, key: str, region: Optional[str] = DEFAULT_CACHE_REGION) -> bool:
|
||||
"""
|
||||
@@ -551,10 +530,7 @@ class AsyncMemoryBackend(AsyncCacheBackend):
|
||||
:param region: 缓存的区
|
||||
:return: 存在返回 True,否则返回 False
|
||||
"""
|
||||
region_cache = self.__get_region_cache(region)
|
||||
if region_cache is None:
|
||||
return False
|
||||
return key in region_cache
|
||||
return self._backend.exists(key=key, region=region)
|
||||
|
||||
async def get(self, key: str, region: Optional[str] = DEFAULT_CACHE_REGION) -> Any:
|
||||
"""
|
||||
@@ -564,10 +540,7 @@ class AsyncMemoryBackend(AsyncCacheBackend):
|
||||
:param region: 缓存的区
|
||||
:return: 返回缓存的值,如果缓存不存在返回 None
|
||||
"""
|
||||
region_cache = self.__get_region_cache(region)
|
||||
if region_cache is None:
|
||||
return None
|
||||
return region_cache.get(key)
|
||||
return self._backend.get(key=key, region=region)
|
||||
|
||||
async def delete(self, key: str, region: Optional[str] = DEFAULT_CACHE_REGION):
|
||||
"""
|
||||
@@ -576,11 +549,7 @@ class AsyncMemoryBackend(AsyncCacheBackend):
|
||||
:param key: 缓存的键
|
||||
:param region: 缓存的区
|
||||
"""
|
||||
region_cache = self.__get_region_cache(region)
|
||||
if region_cache is None:
|
||||
return
|
||||
with lock:
|
||||
del region_cache[key]
|
||||
return self._backend.delete(key=key, region=region)
|
||||
|
||||
async def clear(self, region: Optional[str] = DEFAULT_CACHE_REGION) -> None:
|
||||
"""
|
||||
@@ -588,19 +557,7 @@ class AsyncMemoryBackend(AsyncCacheBackend):
|
||||
|
||||
:param region: 缓存的区,为None时清空所有区缓存
|
||||
"""
|
||||
if region:
|
||||
# 清理指定缓存区
|
||||
region_cache = self.__get_region_cache(region)
|
||||
if region_cache:
|
||||
with lock:
|
||||
region_cache.clear()
|
||||
logger.debug(f"Cleared cache for region: {region}")
|
||||
else:
|
||||
# 清除所有区域的缓存
|
||||
for region_cache in self._region_caches.values():
|
||||
with lock:
|
||||
region_cache.clear()
|
||||
logger.info("All cache cleared!")
|
||||
return self._backend.clear(region=region)
|
||||
|
||||
async def items(self, region: Optional[str] = DEFAULT_CACHE_REGION) -> AsyncGenerator[Tuple[str, Any], None]:
|
||||
"""
|
||||
@@ -609,14 +566,7 @@ class AsyncMemoryBackend(AsyncCacheBackend):
|
||||
:param region: 缓存的区
|
||||
:return: 返回一个字典,包含所有缓存键值对
|
||||
"""
|
||||
region_cache = self.__get_region_cache(region)
|
||||
if region_cache is None:
|
||||
return
|
||||
# 使用锁保护迭代过程,避免在迭代时缓存被修改
|
||||
with lock:
|
||||
# 创建快照避免并发修改问题
|
||||
items_snapshot = list(region_cache.items())
|
||||
for item in items_snapshot:
|
||||
for item in self._backend.items(region):
|
||||
yield item
|
||||
|
||||
async def close(self) -> None:
|
||||
@@ -1115,15 +1065,16 @@ def AsyncCache(cache_type: Literal['ttl', 'lru'] = 'ttl',
|
||||
|
||||
|
||||
def cached(region: Optional[str] = None, maxsize: Optional[int] = 1024, ttl: Optional[int] = None,
|
||||
skip_none: Optional[bool] = True, skip_empty: Optional[bool] = False):
|
||||
skip_none: Optional[bool] = True, skip_empty: Optional[bool] = False, shared_key: Optional[str] = None):
|
||||
"""
|
||||
自定义缓存装饰器,支持为每个 key 动态传递 maxsize 和 ttl
|
||||
|
||||
:param region: 缓存的区
|
||||
:param maxsize: 缓存的最大条目数
|
||||
:param region: 缓存区域的标识符,默认根据模块名、函数名等自动生成标识
|
||||
:param maxsize: 缓存区内的最大条目数
|
||||
:param ttl: 缓存的存活时间,单位秒,未传入则为永久缓存,单位秒
|
||||
:param skip_none: 跳过 None 缓存,默认为 True
|
||||
:param skip_empty: 跳过空值缓存(如 None, [], {}, "", set()),默认为 False
|
||||
:param shared_key: 同步/异步函数共享缓存的键,默认使用函数名(异步函数名会标准化为同步格式,如移除 `async_` 前缀)
|
||||
:return: 装饰器函数
|
||||
"""
|
||||
|
||||
@@ -1173,6 +1124,17 @@ def cached(region: Optional[str] = None, maxsize: Optional[int] = 1024, ttl: Opt
|
||||
return False
|
||||
return True
|
||||
|
||||
def __standardize_func_name() -> str:
|
||||
"""
|
||||
将异步函数名标准化为同步函数的命名,以生成统一的缓存键
|
||||
"""
|
||||
# XXX 假设异步函数名与同步版本仅差`async_`前缀或`_async`后缀(当前MP代码大多符合),否则需通过`shared_key`参数显式指定
|
||||
return (
|
||||
func.__name__.removeprefix("async_").removesuffix("_async")
|
||||
if is_async
|
||||
else func.__name__
|
||||
)
|
||||
|
||||
def __get_cache_key(args, kwargs) -> str:
|
||||
"""
|
||||
根据函数和参数生成缓存键
|
||||
@@ -1194,13 +1156,22 @@ def cached(region: Optional[str] = None, maxsize: Optional[int] = 1024, ttl: Opt
|
||||
bound.arguments[param] for param in signature.parameters if param in bound.arguments
|
||||
]
|
||||
# 使用有序参数生成缓存键
|
||||
return f"{func.__name__}_{hashkey(*keys)}"
|
||||
|
||||
# 获取缓存区
|
||||
cache_region = region if region is not None else f"{func.__module__}.{func.__name__}"
|
||||
return f"{func_name}_{hashkey(*keys)}"
|
||||
|
||||
# 被装饰函数的上层名称(如类名或外层函数名)
|
||||
enclosing_name = (
|
||||
func.__qualname__[:last_dot]
|
||||
if (last_dot := func.__qualname__.rfind(".")) != -1
|
||||
else ""
|
||||
)
|
||||
# 检查是否为异步函数
|
||||
is_async = inspect.iscoroutinefunction(func)
|
||||
# 生成标准化后的函数名称,用于同步/异步函数共享缓存
|
||||
func_name = shared_key if shared_key else __standardize_func_name()
|
||||
# 获取缓存区
|
||||
cache_region = (
|
||||
region if region is not None else f"{func.__module__}:{enclosing_name}:{func_name}"
|
||||
)
|
||||
|
||||
if is_async:
|
||||
# 异步函数使用异步缓存后端
|
||||
|
||||
@@ -209,6 +209,8 @@ class ConfigModel(BaseModel):
|
||||
# ==================== 云盘配置 ====================
|
||||
# 115 AppId
|
||||
U115_APP_ID: str = "100196807"
|
||||
# 115 OAuth2 Server 地址
|
||||
U115_AUTH_SERVER: str = "https://movie-pilot.org"
|
||||
# Alipan AppId
|
||||
ALIPAN_APP_ID: str = "ac1bf04dc9fd4d9aaabb65b4a668d403"
|
||||
|
||||
@@ -219,7 +221,7 @@ class ConfigModel(BaseModel):
|
||||
AUTO_UPDATE_RESOURCE: bool = True
|
||||
|
||||
# ==================== 媒体文件格式配置 ====================
|
||||
# 支持的后缀格式
|
||||
# 支持的视频文件后缀格式
|
||||
RMT_MEDIAEXT: list = Field(
|
||||
default_factory=lambda: ['.mp4', '.mkv', '.ts', '.iso',
|
||||
'.rmvb', '.avi', '.mov', '.mpeg',
|
||||
@@ -230,8 +232,6 @@ class ConfigModel(BaseModel):
|
||||
# 支持的字幕文件后缀格式
|
||||
RMT_SUBEXT: list = Field(default_factory=lambda: ['.srt', '.ass', '.ssa', '.sup'])
|
||||
# 支持的音轨文件后缀格式
|
||||
RMT_AUDIO_TRACK_EXT: list = Field(default_factory=lambda: ['.mka'])
|
||||
# 音轨文件后缀格式
|
||||
RMT_AUDIOEXT: list = Field(
|
||||
default_factory=lambda: ['.aac', '.ac3', '.amr', '.caf', '.cda', '.dsf',
|
||||
'.dff', '.kar', '.m4a', '.mp1', '.mp2', '.mp3',
|
||||
@@ -278,7 +278,7 @@ class ConfigModel(BaseModel):
|
||||
# 搜索多个名称
|
||||
SEARCH_MULTIPLE_NAME: bool = False
|
||||
# 最大搜索名称数量
|
||||
MAX_SEARCH_NAME_LIMIT: int = 2
|
||||
MAX_SEARCH_NAME_LIMIT: int = 3
|
||||
|
||||
# ==================== 下载配置 ====================
|
||||
# 种子标签
|
||||
@@ -305,6 +305,8 @@ class ConfigModel(BaseModel):
|
||||
COOKIECLOUD_BLACKLIST: Optional[str] = None
|
||||
|
||||
# ==================== 整理配置 ====================
|
||||
# 文件整理线程数
|
||||
TRANSFER_THREADS: int = 1
|
||||
# 电影重命名格式
|
||||
MOVIE_RENAME_FORMAT: str = "{{title}}{% if year %} ({{year}}){% endif %}" \
|
||||
"/{{title}}{% if year %} ({{year}}){% endif %}{% if part %}-{{part}}{% endif %}{% if videoFormat %} - {{videoFormat}}{% endif %}" \
|
||||
@@ -320,6 +322,8 @@ class ConfigModel(BaseModel):
|
||||
DEFAULT_SUB: Optional[str] = "zh-cn"
|
||||
# 新增已入库媒体是否跟随TMDB信息变化
|
||||
SCRAP_FOLLOW_TMDB: bool = True
|
||||
# 优先使用辅助识别
|
||||
RECOGNIZE_PLUGIN_FIRST: bool = False
|
||||
|
||||
# ==================== 服务地址配置 ====================
|
||||
# 服务器地址,对应 https://github.com/jxxghp/MoviePilot-Server 项目
|
||||
@@ -337,7 +341,7 @@ class ConfigModel(BaseModel):
|
||||
"https://github.com/thsrite/MoviePilot-Plugins,"
|
||||
"https://github.com/honue/MoviePilot-Plugins,"
|
||||
"https://github.com/InfinityPacer/MoviePilot-Plugins,"
|
||||
"https://github.com/DDS-Derek/MoviePilot-Plugins,"
|
||||
"https://github.com/DDSRem-Dev/MoviePilot-Plugins,"
|
||||
"https://github.com/madrays/MoviePilot-Plugins,"
|
||||
"https://github.com/justzerock/MoviePilot-Plugins,"
|
||||
"https://github.com/KoWming/MoviePilot-Plugins,"
|
||||
@@ -347,7 +351,12 @@ class ConfigModel(BaseModel):
|
||||
"https://github.com/Aqr-K/MoviePilot-Plugins,"
|
||||
"https://github.com/hotlcc/MoviePilot-Plugins-Third,"
|
||||
"https://github.com/gxterry/MoviePilot-Plugins,"
|
||||
"https://github.com/DzAvril/MoviePilot-Plugins")
|
||||
"https://github.com/DzAvril/MoviePilot-Plugins,"
|
||||
"https://github.com/mrtian2016/MoviePilot-Plugins,"
|
||||
"https://github.com/Hqyel/MoviePilot-Plugins-Third,"
|
||||
"https://github.com/xijin285/MoviePilot-Plugins,"
|
||||
"https://github.com/Seed680/MoviePilot-Plugins,"
|
||||
"https://github.com/imaliang/MoviePilot-Plugins")
|
||||
# 插件安装数据共享
|
||||
PLUGIN_STATISTIC_SHARE: bool = True
|
||||
# 是否开启插件热加载
|
||||
@@ -395,6 +404,8 @@ class ConfigModel(BaseModel):
|
||||
SECURITY_IMAGE_SUFFIXES: list = Field(default=[".jpg", ".jpeg", ".png", ".webp", ".gif", ".svg", ".avif"])
|
||||
# PassKey 是否强制用户验证(生物识别等)
|
||||
PASSKEY_REQUIRE_UV: bool = True
|
||||
# 允许在未启用 OTP 时直接注册 PassKey
|
||||
PASSKEY_ALLOW_REGISTER_WITHOUT_OTP: bool = False
|
||||
|
||||
# ==================== 工作流配置 ====================
|
||||
# 工作流数据共享
|
||||
@@ -405,6 +416,8 @@ class ConfigModel(BaseModel):
|
||||
RCLONE_SNAPSHOT_CHECK_FOLDER_MODTIME: bool = True
|
||||
# 对OpenList进行快照对比时,是否检查文件夹的修改时间
|
||||
OPENLIST_SNAPSHOT_CHECK_FOLDER_MODTIME: bool = True
|
||||
# 对阿里云盘进行快照对比时,是否检查文件夹的修改时间(默认关闭,因为阿里云盘目录时间不随子文件变更而更新)
|
||||
ALIPAN_SNAPSHOT_CHECK_FOLDER_MODTIME: bool = False
|
||||
|
||||
# ==================== Docker配置 ====================
|
||||
# Docker Client API地址
|
||||
@@ -425,10 +438,12 @@ class ConfigModel(BaseModel):
|
||||
LLM_API_KEY: Optional[str] = None
|
||||
# LLM基础URL(用于自定义API端点)
|
||||
LLM_BASE_URL: Optional[str] = "https://api.deepseek.com"
|
||||
# LLM最大上下文Token数量(K)
|
||||
LLM_MAX_CONTEXT_TOKENS: int = 64
|
||||
# LLM温度参数
|
||||
LLM_TEMPERATURE: float = 0.1
|
||||
# LLM最大迭代次数
|
||||
LLM_MAX_ITERATIONS: int = 15
|
||||
LLM_MAX_ITERATIONS: int = 128
|
||||
# LLM工具调用超时时间(秒)
|
||||
LLM_TOOL_TIMEOUT: int = 300
|
||||
# 是否启用详细日志
|
||||
@@ -439,6 +454,16 @@ class ConfigModel(BaseModel):
|
||||
LLM_MEMORY_RETENTION_DAYS: int = 1
|
||||
# Redis记忆保留天数(如果使用Redis)
|
||||
LLM_REDIS_MEMORY_RETENTION_DAYS: int = 7
|
||||
# 是否启用AI推荐
|
||||
AI_RECOMMEND_ENABLED: bool = False
|
||||
# AI推荐用户偏好
|
||||
AI_RECOMMEND_USER_PREFERENCE: str = ""
|
||||
# Tavily API密钥(用于网络搜索)
|
||||
TAVILY_API_KEY: str = "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"
|
||||
|
||||
# AI推荐条目数量限制
|
||||
AI_RECOMMEND_MAX_ITEMS: int = 50
|
||||
|
||||
|
||||
|
||||
class Settings(BaseSettings, ConfigModel, LogConfigModel):
|
||||
@@ -843,6 +868,22 @@ class Settings(BaseSettings, ConfigModel, LogConfigModel):
|
||||
rename_format = re.sub(r'/+', '/', rename_format)
|
||||
return rename_format.strip("/")
|
||||
|
||||
def TMDB_IMAGE_URL(
|
||||
self, file_path: Optional[str], file_size: str = "original"
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
获取TMDB图片网址
|
||||
|
||||
:param file_path: TMDB API返回的xxx_path
|
||||
:param file_size: 图片大小,例如:'original', 'w500' 等
|
||||
:return: 图片的完整URL,如果 file_path 为空则返回 None
|
||||
"""
|
||||
if not file_path:
|
||||
return None
|
||||
return (
|
||||
f"https://{self.TMDB_IMAGE_DOMAIN}/t/p/{file_size}/{file_path.removeprefix('/')}"
|
||||
)
|
||||
|
||||
|
||||
# 实例化配置
|
||||
settings = Settings()
|
||||
|
||||
@@ -465,7 +465,7 @@ class MediaInfo:
|
||||
for seainfo in info.get('seasons'):
|
||||
# 季
|
||||
season = seainfo.get("season_number")
|
||||
if not season:
|
||||
if season is None:
|
||||
continue
|
||||
# 集
|
||||
episode_count = seainfo.get("episode_count")
|
||||
@@ -479,11 +479,11 @@ class MediaInfo:
|
||||
self.episode_groups = info.pop("episode_groups").get("results") or []
|
||||
|
||||
# 海报
|
||||
if info.get('poster_path'):
|
||||
self.poster_path = f"https://{settings.TMDB_IMAGE_DOMAIN}/t/p/original{info.get('poster_path')}"
|
||||
if path := info.get('poster_path'):
|
||||
self.poster_path = settings.TMDB_IMAGE_URL(path)
|
||||
# 背景
|
||||
if info.get('backdrop_path'):
|
||||
self.backdrop_path = f"https://{settings.TMDB_IMAGE_DOMAIN}/t/p/original{info.get('backdrop_path')}"
|
||||
if path := info.get('backdrop_path'):
|
||||
self.backdrop_path = settings.TMDB_IMAGE_URL(path)
|
||||
# 导演和演员
|
||||
self.directors, self.actors = __directors_actors(info)
|
||||
# 别名和译名
|
||||
@@ -545,9 +545,9 @@ class MediaInfo:
|
||||
# 识别标题中的季
|
||||
meta = MetaInfo(info.get("title"))
|
||||
# 季
|
||||
if not self.season:
|
||||
if self.season is None:
|
||||
self.season = meta.begin_season
|
||||
if self.season:
|
||||
if self.season is not None:
|
||||
self.type = MediaType.TV
|
||||
elif not self.type:
|
||||
self.type = MediaType.MOVIE
|
||||
@@ -607,13 +607,13 @@ class MediaInfo:
|
||||
# 剧集
|
||||
if self.type == MediaType.TV and not self.seasons:
|
||||
meta = MetaInfo(info.get("title"))
|
||||
season = meta.begin_season or 1
|
||||
season = meta.begin_season if meta.begin_season is not None else 1
|
||||
episodes_count = info.get("episodes_count")
|
||||
if episodes_count:
|
||||
self.seasons[season] = list(range(1, episodes_count + 1))
|
||||
# 季年份
|
||||
if self.type == MediaType.TV and not self.season_years:
|
||||
season = self.season or 1
|
||||
season = self.season if self.season is not None else 1
|
||||
self.season_years = {
|
||||
season: self.year
|
||||
}
|
||||
@@ -667,7 +667,7 @@ class MediaInfo:
|
||||
# 识别标题中的季
|
||||
meta = MetaInfo(self.title)
|
||||
# 季
|
||||
if not self.season:
|
||||
if self.season is None:
|
||||
self.season = meta.begin_season
|
||||
# 评分
|
||||
if not self.vote_average:
|
||||
@@ -703,7 +703,7 @@ class MediaInfo:
|
||||
# 剧集
|
||||
if self.type == MediaType.TV and not self.seasons:
|
||||
meta = MetaInfo(self.title)
|
||||
season = meta.begin_season or 1
|
||||
season = meta.begin_season if meta.begin_season is not None else 1
|
||||
episodes_count = info.get("total_episodes")
|
||||
if episodes_count:
|
||||
self.seasons[season] = list(range(1, episodes_count + 1))
|
||||
|
||||
@@ -17,6 +17,7 @@ class MetaAnime(MetaBase):
|
||||
"""
|
||||
_anime_no_words = ['CHS&CHT', 'MP4', 'GB MP4', 'WEB-DL']
|
||||
_name_nostring_re = r"S\d{2}\s*-\s*S\d{2}|S\d{2}|\s+S\d{1,2}|EP?\d{2,4}\s*-\s*EP?\d{2,4}|EP?\d{2,4}|\s+EP?\d{1,4}|\s+GB"
|
||||
_fps_re = r"(\d{2,3})(?=FPS)"
|
||||
|
||||
def __init__(self, title: str, subtitle: str = None, isfile: bool = False):
|
||||
super().__init__(title, subtitle, isfile)
|
||||
@@ -173,6 +174,8 @@ class MetaAnime(MetaBase):
|
||||
self.audio_encode = anitopy_info.get("audio_term")
|
||||
if isinstance(self.audio_encode, list):
|
||||
self.audio_encode = self.audio_encode[0]
|
||||
# 帧率信息
|
||||
self.__init_anime_fps(anitopy_info, original_title)
|
||||
# 解析副标题,只要季和集
|
||||
self.init_subtitle(self.org_string)
|
||||
if not self._subtitle_flag and self.subtitle:
|
||||
@@ -182,6 +185,20 @@ class MetaAnime(MetaBase):
|
||||
except Exception as e:
|
||||
logger.error(f"解析动漫信息失败:{str(e)} - {traceback.format_exc()}")
|
||||
|
||||
def __init_anime_fps(self, anitopy_info: dict, original_title: str):
|
||||
"""
|
||||
从原始标题中提取帧率信息,与MetaVideo保持完全一致的实现
|
||||
"""
|
||||
re_res = re.search(rf"({self._fps_re})", original_title, re.IGNORECASE)
|
||||
if re_res:
|
||||
fps_value = None
|
||||
if re_res.group(1): # FPS格式
|
||||
fps_value = re_res.group(1)
|
||||
|
||||
if fps_value and fps_value.isdigit():
|
||||
# 只存储纯数值
|
||||
self.fps = int(fps_value)
|
||||
|
||||
@staticmethod
|
||||
def __prepare_title(title: str):
|
||||
"""
|
||||
|
||||
@@ -66,6 +66,9 @@ class MetaBase(object):
|
||||
# 附加信息
|
||||
tmdbid: int = None
|
||||
doubanid: str = None
|
||||
# 帧率信息(纯数值)
|
||||
fps: Optional[int] = None
|
||||
|
||||
|
||||
# 副标题解析
|
||||
_subtitle_flag = False
|
||||
@@ -448,6 +451,13 @@ class MetaBase(object):
|
||||
"""
|
||||
return self.audio_encode or ""
|
||||
|
||||
@property
|
||||
def frame_rate(self) -> int:
|
||||
"""
|
||||
返回帧率信息
|
||||
"""
|
||||
return self.fps or None
|
||||
|
||||
def is_in_season(self, season: Union[list, int, str]) -> bool:
|
||||
"""
|
||||
是否包含季
|
||||
@@ -535,7 +545,7 @@ class MetaBase(object):
|
||||
|
||||
def merge(self, meta: Self):
|
||||
"""
|
||||
全并Meta信息
|
||||
合并Meta信息
|
||||
"""
|
||||
# 类型
|
||||
if self.type == MediaType.UNKNOWN \
|
||||
@@ -581,6 +591,9 @@ class MetaBase(object):
|
||||
# 音频编码
|
||||
if not self.audio_encode:
|
||||
self.audio_encode = meta.audio_encode
|
||||
# 帧率信息
|
||||
if not self.fps:
|
||||
self.fps = meta.fps
|
||||
# Part
|
||||
if not self.part:
|
||||
self.part = meta.part
|
||||
|
||||
@@ -53,7 +53,7 @@ class MetaVideo(MetaBase):
|
||||
_resources_pix_re2 = r"(^[248]+K)"
|
||||
_video_encode_re = r"^(H26[45])$|^(x26[45])$|^AVC$|^HEVC$|^VC\d?$|^MPEG\d?$|^Xvid$|^DivX$|^AV1$|^HDR\d*$|^AVS(\+|[23])$"
|
||||
_audio_encode_re = r"^DTS\d?$|^DTSHD$|^DTSHDMA$|^Atmos$|^TrueHD\d?$|^AC3$|^\dAudios?$|^DDP\d?$|^DD\+\d?$|^DD\d?$|^LPCM\d?$|^AAC\d?$|^FLAC\d?$|^HD\d?$|^MA\d?$|^HR\d?$|^Opus\d?$|^Vorbis\d?$|^AV[3S]A$"
|
||||
|
||||
_fps_re = r"(\d{2,3})(?=FPS)"
|
||||
def __init__(self, title: str, subtitle: str = None, isfile: bool = False):
|
||||
"""
|
||||
初始化
|
||||
@@ -76,7 +76,7 @@ class MetaVideo(MetaBase):
|
||||
self.type = MediaType.TV
|
||||
return
|
||||
# 全名为Season xx 及 Sxx 直接返回
|
||||
season_full_res = re.search(r"^Season\s+(\d{1,3})$|^S(\d{1,3})$", title)
|
||||
season_full_res = re.search(r"^(?:Season\s+|S)(\d{1,3})$", title, re.IGNORECASE)
|
||||
if season_full_res:
|
||||
self.type = MediaType.TV
|
||||
season = season_full_res.group(1)
|
||||
@@ -129,6 +129,9 @@ class MetaVideo(MetaBase):
|
||||
# 音频编码
|
||||
if self._continue_flag:
|
||||
self.__init_audio_encode(token)
|
||||
# 帧率
|
||||
if self._continue_flag:
|
||||
self.__init_fps(token)
|
||||
# 取下一个,直到没有为卡
|
||||
token = tokens.get_next()
|
||||
self._continue_flag = True
|
||||
@@ -301,7 +304,8 @@ class MetaVideo(MetaBase):
|
||||
return
|
||||
else:
|
||||
# 后缀名不要
|
||||
if ".%s".lower() % token in settings.RMT_MEDIAEXT:
|
||||
media_exts = settings.RMT_MEDIAEXT + settings.RMT_SUBEXT + settings.RMT_AUDIOEXT
|
||||
if ".%s".lower() % token in media_exts:
|
||||
return
|
||||
# 英文或者英文+数字,拼装起来
|
||||
if self.en_name:
|
||||
@@ -715,3 +719,25 @@ class MetaVideo(MetaBase):
|
||||
else:
|
||||
self.audio_encode = "%s %s" % (self.audio_encode, token)
|
||||
self._last_token = token
|
||||
|
||||
def __init_fps(self, token: str):
|
||||
"""
|
||||
识别帧率
|
||||
"""
|
||||
if not self.name:
|
||||
return
|
||||
|
||||
re_res = re.search(rf"({self._fps_re})", token, re.IGNORECASE)
|
||||
if re_res:
|
||||
self._continue_flag = False
|
||||
self._stop_name_flag = True
|
||||
self._last_token_type = "fps"
|
||||
# 提取帧率数值
|
||||
fps_value = None
|
||||
if re_res.group(1): # FPS格式
|
||||
fps_value = re_res.group(1)
|
||||
|
||||
if fps_value and fps_value.isdigit():
|
||||
# 只存储纯数值
|
||||
self.fps = int(fps_value)
|
||||
self._last_token = f"{self.fps}FPS"
|
||||
|
||||
@@ -52,6 +52,7 @@ class ReleaseGroupsMatcher(metaclass=Singleton):
|
||||
"nicept": [],
|
||||
"oshen": [],
|
||||
"ourbits": ['Our(?:Bits|TV)', 'FLTTH', 'Ao', 'PbK', 'MGs', 'iLove(?:HD|TV)'],
|
||||
"panda": ['Panda', 'AilMWeb'],
|
||||
"piggo": ['PiGo(?:NF|(?:H|WE)B)'],
|
||||
"ptchina": [],
|
||||
"pterclub": ['PTer(?:DIY|Game|(?:M|T)V|WEB|)'],
|
||||
@@ -105,7 +106,7 @@ class ReleaseGroupsMatcher(metaclass=Singleton):
|
||||
else:
|
||||
groups = self.__release_groups
|
||||
title = f"{title} "
|
||||
groups_re = re.compile(r"(?<=[-@\[£【&])(?:(?:%s))(?=[@.\s\S\]\[】&])" % groups, re.I)
|
||||
groups_re = re.compile(r"(?<=[-@\[£【&])(?:(?:%s))(?=$|[@.\s\]\[】&])" % groups, re.I)
|
||||
unique_groups = []
|
||||
for item in re.findall(groups_re, title):
|
||||
item_str = item[0] if isinstance(item, tuple) else item
|
||||
|
||||
@@ -25,7 +25,8 @@ def MetaInfo(title: str, subtitle: Optional[str] = None, custom_words: List[str]
|
||||
# 获取标题中媒体信息
|
||||
title, metainfo = find_metainfo(title)
|
||||
# 判断是否处理文件
|
||||
if title and Path(title).suffix.lower() in settings.RMT_MEDIAEXT:
|
||||
media_exts = settings.RMT_MEDIAEXT + settings.RMT_SUBEXT + settings.RMT_AUDIOEXT
|
||||
if title and Path(title).suffix.lower() in media_exts:
|
||||
isfile = True
|
||||
# 去掉后缀
|
||||
title = Path(title).stem
|
||||
@@ -62,21 +63,24 @@ def MetaInfo(title: str, subtitle: Optional[str] = None, custom_words: List[str]
|
||||
return meta
|
||||
|
||||
|
||||
def MetaInfoPath(path: Path) -> MetaBase:
|
||||
def MetaInfoPath(path: Path, custom_words: List[str] = None) -> MetaBase:
|
||||
"""
|
||||
根据路径识别元数据
|
||||
:param path: 路径
|
||||
:param custom_words: 自定义识别词列表
|
||||
"""
|
||||
# 文件元数据,不包含后缀
|
||||
file_meta = MetaInfo(title=path.name)
|
||||
file_meta = MetaInfo(title=path.name, custom_words=custom_words)
|
||||
# 上级目录元数据
|
||||
dir_meta = MetaInfo(title=path.parent.name)
|
||||
# 合并元数据
|
||||
file_meta.merge(dir_meta)
|
||||
dir_meta = MetaInfo(title=path.parent.name, custom_words=custom_words)
|
||||
if file_meta.type == MediaType.TV or dir_meta.type != MediaType.TV:
|
||||
# 合并元数据
|
||||
file_meta.merge(dir_meta)
|
||||
# 上上级目录元数据
|
||||
root_meta = MetaInfo(title=path.parent.parent.name)
|
||||
# 合并元数据
|
||||
file_meta.merge(root_meta)
|
||||
root_meta = MetaInfo(title=path.parent.parent.name, custom_words=custom_words)
|
||||
if file_meta.type == MediaType.TV or root_meta.type != MediaType.TV:
|
||||
# 合并元数据
|
||||
file_meta.merge(root_meta)
|
||||
return file_meta
|
||||
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ import concurrent.futures
|
||||
import importlib.util
|
||||
import inspect
|
||||
import os
|
||||
import posixpath
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
@@ -775,11 +776,17 @@ class PluginManager(ConfigReloadMixin, metaclass=Singleton):
|
||||
:param dist_path: 插件的分发路径
|
||||
:return: 远程入口地址
|
||||
"""
|
||||
if dist_path.startswith("/"):
|
||||
dist_path = dist_path[1:]
|
||||
if dist_path.endswith("/"):
|
||||
dist_path = dist_path[:-1]
|
||||
return f"/plugin/file/{plugin_id.lower()}/{dist_path}/remoteEntry.js"
|
||||
dist_path = dist_path.strip("/")
|
||||
path = posixpath.join(
|
||||
"plugin",
|
||||
"file",
|
||||
plugin_id.lower(),
|
||||
dist_path,
|
||||
"remoteEntry.js",
|
||||
)
|
||||
if not path.startswith("/"):
|
||||
path = "/" + path
|
||||
return path
|
||||
|
||||
def get_plugin_remotes(self, pid: Optional[str] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
|
||||
@@ -17,6 +17,7 @@ from fastapi.security import OAuth2PasswordBearer, APIKeyHeader, APIKeyQuery, AP
|
||||
from passlib.context import CryptContext
|
||||
|
||||
from app import schemas
|
||||
from app.core.cache import cached
|
||||
from app.core.config import settings
|
||||
from app.log import logger
|
||||
|
||||
@@ -24,7 +25,8 @@ pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
||||
ALGORITHM = "HS256"
|
||||
|
||||
# OAuth2PasswordBearer 用于 JWT Token 认证
|
||||
oauth2_scheme = OAuth2PasswordBearer(
|
||||
oauth2_scheme_manual_error = OAuth2PasswordBearer(
|
||||
auto_error=False, # 禁用自动错误处理,用以支持API令牌鉴权
|
||||
tokenUrl=f"{settings.API_V1_STR}/login/access-token"
|
||||
)
|
||||
|
||||
@@ -41,6 +43,58 @@ api_key_header = APIKeyHeader(name="X-API-KEY", auto_error=False, scheme_name="a
|
||||
api_key_query = APIKeyQuery(name="apikey", auto_error=False, scheme_name="api_key_query")
|
||||
|
||||
|
||||
def __get_api_token(
|
||||
token_query: Annotated[str | None, Security(api_token_query)] = None
|
||||
) -> str | None:
|
||||
"""
|
||||
从 URL 查询参数中获取 API Token
|
||||
:param token_query: 从 URL 中的 `token` 查询参数获取 API Token
|
||||
:return: 返回获取到的 API Token,若无则返回 None
|
||||
"""
|
||||
return token_query
|
||||
|
||||
|
||||
def __get_api_key(
|
||||
key_query: Annotated[str | None, Security(api_key_query)] = None,
|
||||
key_header: Annotated[str | None, Security(api_key_header)] = None
|
||||
) -> str | None:
|
||||
"""
|
||||
从 URL 查询参数或请求头部获取 API Key,优先使用请求头
|
||||
:param key_query: URL 中的 `apikey` 查询参数
|
||||
:param key_header: 请求头中的 `X-API-KEY` 参数
|
||||
:return: 返回从 URL 或请求头中获取的 API Key,若无则返回 None
|
||||
"""
|
||||
return key_header or key_query # 首选请求头
|
||||
|
||||
|
||||
@cached(maxsize=1, ttl=600)
|
||||
def __create_superuser_token_payload() -> schemas.TokenPayload:
|
||||
"""
|
||||
创建管理员用户的TokenPayload
|
||||
|
||||
:return: 管理员TokenPayload
|
||||
"""
|
||||
# 延迟导入
|
||||
# pylint: disable=import-outside-toplevel
|
||||
# pylint: disable=no-name-in-module
|
||||
from app.db.user_oper import UserOper
|
||||
from app.helper.sites import SitesHelper # noqa
|
||||
|
||||
user = UserOper().get_by_name(settings.SUPERUSER)
|
||||
if not user or not user.is_superuser:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="用户权限不足",
|
||||
)
|
||||
return schemas.TokenPayload(
|
||||
sub=user.id,
|
||||
username=user.name,
|
||||
super_user=user.is_superuser,
|
||||
level=SitesHelper().auth_level,
|
||||
purpose="authentication",
|
||||
)
|
||||
|
||||
|
||||
def create_access_token(
|
||||
userid: Union[str, Any],
|
||||
username: str,
|
||||
@@ -176,23 +230,43 @@ def __verify_token(token: str, purpose: Optional[str] = "authentication") -> sch
|
||||
def verify_token(
|
||||
request: Request,
|
||||
response: Response,
|
||||
token: Annotated[str, Security(oauth2_scheme)]
|
||||
jwt_token: Annotated[str | None, Security(oauth2_scheme_manual_error)],
|
||||
api_key: Annotated[str | None, Security(__get_api_key)],
|
||||
api_token: Annotated[str | None, Security(__get_api_token)],
|
||||
) -> schemas.TokenPayload:
|
||||
"""
|
||||
验证 JWT 令牌并自动处理 resource_token 写入
|
||||
|
||||
如果缺少JWT令牌再尝试用API令牌鉴权
|
||||
|
||||
:param request: 请求对象,用于访问 Cookie 和请求信息
|
||||
:param response: 响应对象,用于设置 Cookie
|
||||
:param token: 从 Authorization 头部获取的 JWT 令牌
|
||||
:param jwt_token: 从 Authorization 头部获取的 JWT 令牌
|
||||
:param api_key: 从 查询参数`apikey` 或 请求头`X-API-KEY` 获取 API Token
|
||||
:param api_token: 从 查询参数`token` 获取 API Token
|
||||
:return: 解析后的 TokenPayload
|
||||
:raises HTTPException: 如果令牌无效或用途不匹配
|
||||
"""
|
||||
# 验证并解析 JWT 认证令牌
|
||||
payload = __verify_token(token=token, purpose="authentication")
|
||||
if jwt_token:
|
||||
# 验证并解析 JWT 认证令牌
|
||||
payload = __verify_token(token=jwt_token, purpose="authentication")
|
||||
|
||||
# 如果没有 resource_token,生成并写入到 Cookie
|
||||
__set_or_refresh_resource_token_cookie(request, response, payload)
|
||||
# 如果没有 resource_token,生成并写入到 Cookie
|
||||
__set_or_refresh_resource_token_cookie(request, response, payload)
|
||||
|
||||
return payload
|
||||
return payload
|
||||
elif api_key:
|
||||
verify_apikey(api_key)
|
||||
return __create_superuser_token_payload()
|
||||
elif api_token:
|
||||
verify_apitoken(api_token)
|
||||
return __create_superuser_token_payload()
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Not authenticated",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
|
||||
def verify_resource_token(
|
||||
@@ -208,31 +282,7 @@ def verify_resource_token(
|
||||
return __verify_token(token=resource_token, purpose="resource")
|
||||
|
||||
|
||||
def __get_api_token(
|
||||
token_query: Annotated[str | None, Security(api_token_query)] = None
|
||||
) -> str:
|
||||
"""
|
||||
从 URL 查询参数中获取 API Token
|
||||
:param token_query: 从 URL 中的 `token` 查询参数获取 API Token
|
||||
:return: 返回获取到的 API Token,若无则返回 None
|
||||
"""
|
||||
return token_query
|
||||
|
||||
|
||||
def __get_api_key(
|
||||
key_query: Annotated[str | None, Security(api_key_query)] = None,
|
||||
key_header: Annotated[str | None, Security(api_key_header)] = None
|
||||
) -> str:
|
||||
"""
|
||||
从 URL 查询参数或请求头部获取 API Key,优先使用 URL 参数
|
||||
:param key_query: URL 中的 `apikey` 查询参数
|
||||
:param key_header: 请求头中的 `X-API-KEY` 参数
|
||||
:return: 返回从 URL 或请求头中获取的 API Key,若无则返回 None
|
||||
"""
|
||||
return key_query or key_header
|
||||
|
||||
|
||||
def __verify_key(key: str, expected_key: str, key_type: str) -> str:
|
||||
def __verify_key(key: str | None, expected_key: str, key_type: str) -> str:
|
||||
"""
|
||||
通用的 API Key 或 Token 验证函数
|
||||
:param key: 从请求中获取的 API Key 或 Token
|
||||
@@ -241,7 +291,7 @@ def __verify_key(key: str, expected_key: str, key_type: str) -> str:
|
||||
:return: 返回校验通过的 API Key 或 Token
|
||||
:raises HTTPException: 如果校验不通过,抛出 401 错误
|
||||
"""
|
||||
if key != expected_key:
|
||||
if not key or key != expected_key:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=f"{key_type} 校验不通过"
|
||||
@@ -249,7 +299,7 @@ def __verify_key(key: str, expected_key: str, key_type: str) -> str:
|
||||
return key
|
||||
|
||||
|
||||
def verify_apitoken(token: Annotated[str, Security(__get_api_token)]) -> str:
|
||||
def verify_apitoken(token: Annotated[str | None, Security(__get_api_token)]) -> str:
|
||||
"""
|
||||
使用 API Token 进行身份认证
|
||||
:param token: API Token,从 URL 查询参数中获取 token=xxx
|
||||
@@ -258,7 +308,7 @@ def verify_apitoken(token: Annotated[str, Security(__get_api_token)]) -> str:
|
||||
return __verify_key(token, settings.API_TOKEN, "token")
|
||||
|
||||
|
||||
def verify_apikey(apikey: Annotated[str, Security(__get_api_key)]) -> str:
|
||||
def verify_apikey(apikey: Annotated[str | None, Security(__get_api_key)]) -> str:
|
||||
"""
|
||||
使用 API Key 进行身份认证
|
||||
:param apikey: API Key,从 URL 查询参数中获取 apikey=xxx,或请求头中获取 X-API-KEY=xxx
|
||||
|
||||
@@ -454,7 +454,6 @@ class Base:
|
||||
|
||||
@db_update
|
||||
def update(self, db: Session, payload: dict):
|
||||
payload = {k: v for k, v in payload.items() if v is not None}
|
||||
for key, value in payload.items():
|
||||
setattr(self, key, value)
|
||||
if inspect(self).detached:
|
||||
@@ -462,7 +461,6 @@ class Base:
|
||||
|
||||
@async_db_update
|
||||
async def async_update(self, db: AsyncSession, payload: dict):
|
||||
payload = {k: v for k, v in payload.items() if v is not None}
|
||||
for key, value in payload.items():
|
||||
setattr(self, key, value)
|
||||
if inspect(self).detached:
|
||||
|
||||
@@ -49,7 +49,7 @@ class MediaServerOper(DbOper):
|
||||
if not item:
|
||||
return None
|
||||
|
||||
if kwargs.get("season"):
|
||||
if kwargs.get("season") is not None:
|
||||
# 判断季是否存在
|
||||
if not item.seasoninfo:
|
||||
return None
|
||||
@@ -75,7 +75,7 @@ class MediaServerOper(DbOper):
|
||||
if not item:
|
||||
return None
|
||||
|
||||
if kwargs.get("season"):
|
||||
if kwargs.get("season") is not None:
|
||||
# 判断季是否存在
|
||||
if not item.seasoninfo:
|
||||
return None
|
||||
|
||||
@@ -55,6 +55,8 @@ class DownloadHistory(Base):
|
||||
media_category = Column(String)
|
||||
# 剧集组
|
||||
episode_group = Column(String)
|
||||
# 自定义识别词(用于整理时应用)
|
||||
custom_words = Column(String)
|
||||
|
||||
@classmethod
|
||||
@db_query
|
||||
@@ -102,14 +104,14 @@ class DownloadHistory(Base):
|
||||
# TMDBID + 类型
|
||||
if tmdbid and mtype:
|
||||
# 电视剧某季某集
|
||||
if season and episode:
|
||||
if season is not None and episode:
|
||||
return db.query(DownloadHistory).filter(DownloadHistory.tmdbid == tmdbid,
|
||||
DownloadHistory.type == mtype,
|
||||
DownloadHistory.seasons == season,
|
||||
DownloadHistory.episodes == episode).order_by(
|
||||
DownloadHistory.id.desc()).all()
|
||||
# 电视剧某季
|
||||
elif season:
|
||||
elif season is not None:
|
||||
return db.query(DownloadHistory).filter(DownloadHistory.tmdbid == tmdbid,
|
||||
DownloadHistory.type == mtype,
|
||||
DownloadHistory.seasons == season).order_by(
|
||||
@@ -122,14 +124,14 @@ class DownloadHistory(Base):
|
||||
# 标题 + 年份
|
||||
elif title and year:
|
||||
# 电视剧某季某集
|
||||
if season and episode:
|
||||
if season is not None and episode:
|
||||
return db.query(DownloadHistory).filter(DownloadHistory.title == title,
|
||||
DownloadHistory.year == year,
|
||||
DownloadHistory.seasons == season,
|
||||
DownloadHistory.episodes == episode).order_by(
|
||||
DownloadHistory.id.desc()).all()
|
||||
# 电视剧某季
|
||||
elif season:
|
||||
elif season is not None:
|
||||
return db.query(DownloadHistory).filter(DownloadHistory.title == title,
|
||||
DownloadHistory.year == year,
|
||||
DownloadHistory.seasons == season).order_by(
|
||||
@@ -207,7 +209,7 @@ class DownloadFiles(Base):
|
||||
@classmethod
|
||||
@db_query
|
||||
def get_by_hash(cls, db: Session, download_hash: str, state: Optional[int] = None):
|
||||
if state:
|
||||
if state is not None:
|
||||
return db.query(cls).filter(cls.download_hash == download_hash,
|
||||
cls.state == state).all()
|
||||
else:
|
||||
|
||||
@@ -93,7 +93,7 @@ class Subscribe(Base):
|
||||
def exists(cls, db: Session, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||||
season: Optional[int] = None):
|
||||
if tmdbid:
|
||||
if season:
|
||||
if season is not None:
|
||||
return db.query(cls).filter(cls.tmdbid == tmdbid,
|
||||
cls.season == season).first()
|
||||
return db.query(cls).filter(cls.tmdbid == tmdbid).first()
|
||||
@@ -106,7 +106,7 @@ class Subscribe(Base):
|
||||
async def async_exists(cls, db: AsyncSession, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||||
season: Optional[int] = None):
|
||||
if tmdbid:
|
||||
if season:
|
||||
if season is not None:
|
||||
result = await db.execute(
|
||||
select(cls).filter(cls.tmdbid == tmdbid, cls.season == season)
|
||||
)
|
||||
@@ -148,7 +148,7 @@ class Subscribe(Base):
|
||||
@classmethod
|
||||
@db_query
|
||||
def get_by_title(cls, db: Session, title: str, season: Optional[int] = None):
|
||||
if season:
|
||||
if season is not None:
|
||||
return db.query(cls).filter(cls.name == title,
|
||||
cls.season == season).first()
|
||||
return db.query(cls).filter(cls.name == title).first()
|
||||
@@ -156,7 +156,7 @@ class Subscribe(Base):
|
||||
@classmethod
|
||||
@async_db_query
|
||||
async def async_get_by_title(cls, db: AsyncSession, title: str, season: Optional[int] = None):
|
||||
if season:
|
||||
if season is not None:
|
||||
result = await db.execute(
|
||||
select(cls).filter(cls.name == title, cls.season == season)
|
||||
)
|
||||
@@ -169,7 +169,7 @@ class Subscribe(Base):
|
||||
@classmethod
|
||||
@db_query
|
||||
def get_by_tmdbid(cls, db: Session, tmdbid: int, season: Optional[int] = None):
|
||||
if season:
|
||||
if season is not None:
|
||||
return db.query(cls).filter(cls.tmdbid == tmdbid,
|
||||
cls.season == season).all()
|
||||
else:
|
||||
@@ -178,7 +178,7 @@ class Subscribe(Base):
|
||||
@classmethod
|
||||
@async_db_query
|
||||
async def async_get_by_tmdbid(cls, db: AsyncSession, tmdbid: int, season: Optional[int] = None):
|
||||
if season:
|
||||
if season is not None:
|
||||
result = await db.execute(
|
||||
select(cls).filter(cls.tmdbid == tmdbid, cls.season == season)
|
||||
)
|
||||
@@ -227,6 +227,66 @@ class Subscribe(Base):
|
||||
)
|
||||
return result.scalars().first()
|
||||
|
||||
@classmethod
|
||||
@db_query
|
||||
def get_by(cls, db: Session, type: str, season: Optional[str] = None,
|
||||
tmdbid: Optional[int] = None, doubanid: Optional[str] = None, bangumiid: Optional[str] = None):
|
||||
"""
|
||||
根据条件查询订阅
|
||||
"""
|
||||
# TMDBID
|
||||
if tmdbid:
|
||||
if season is not None:
|
||||
result = db.query(cls).filter(
|
||||
cls.tmdbid == tmdbid, cls.type == type, cls.season == season
|
||||
)
|
||||
else:
|
||||
result = db.query(cls).filter(cls.tmdbid == tmdbid, cls.type == type)
|
||||
# 豆瓣ID
|
||||
elif doubanid:
|
||||
result = db.query(cls).filter(cls.doubanid == doubanid, cls.type == type)
|
||||
# BangumiID
|
||||
elif bangumiid:
|
||||
result = db.query(cls).filter(cls.bangumiid == bangumiid, cls.type == type)
|
||||
else:
|
||||
return None
|
||||
|
||||
return result.first()
|
||||
|
||||
@classmethod
|
||||
@async_db_query
|
||||
async def async_get_by(cls, db: AsyncSession, type: str, season: Optional[str] = None,
|
||||
tmdbid: Optional[int] = None, doubanid: Optional[str] = None, bangumiid: Optional[str] = None):
|
||||
"""
|
||||
根据条件查询订阅
|
||||
"""
|
||||
# TMDBID
|
||||
if tmdbid:
|
||||
if season is not None:
|
||||
result = await db.execute(
|
||||
select(cls).filter(
|
||||
cls.tmdbid == tmdbid, cls.type == type, cls.season == season
|
||||
)
|
||||
)
|
||||
else:
|
||||
result = await db.execute(
|
||||
select(cls).filter(cls.tmdbid == tmdbid, cls.type == type)
|
||||
)
|
||||
# 豆瓣ID
|
||||
elif doubanid:
|
||||
result = await db.execute(
|
||||
select(cls).filter(cls.doubanid == doubanid, cls.type == type)
|
||||
)
|
||||
# BangumiID
|
||||
elif bangumiid:
|
||||
result = await db.execute(
|
||||
select(cls).filter(cls.bangumiid == bangumiid, cls.type == type)
|
||||
)
|
||||
else:
|
||||
return None
|
||||
|
||||
return result.scalars().first()
|
||||
|
||||
@db_update
|
||||
def delete_by_tmdbid(self, db: Session, tmdbid: int, season: int):
|
||||
subscrbies = self.get_by_tmdbid(db, tmdbid, season)
|
||||
|
||||
@@ -99,7 +99,7 @@ class SubscribeHistory(Base):
|
||||
def exists(cls, db: Session, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||||
season: Optional[int] = None):
|
||||
if tmdbid:
|
||||
if season:
|
||||
if season is not None:
|
||||
return db.query(cls).filter(cls.tmdbid == tmdbid,
|
||||
cls.season == season).first()
|
||||
return db.query(cls).filter(cls.tmdbid == tmdbid).first()
|
||||
@@ -112,7 +112,7 @@ class SubscribeHistory(Base):
|
||||
async def async_exists(cls, db: AsyncSession, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||||
season: Optional[int] = None):
|
||||
if tmdbid:
|
||||
if season:
|
||||
if season is not None:
|
||||
result = await db.execute(
|
||||
select(cls).filter(cls.tmdbid == tmdbid, cls.season == season)
|
||||
)
|
||||
|
||||
@@ -266,14 +266,14 @@ class TransferHistory(Base):
|
||||
# TMDBID + 类型
|
||||
if tmdbid and mtype:
|
||||
# 电视剧某季某集
|
||||
if season and episode:
|
||||
if season is not None and episode:
|
||||
return db.query(cls).filter(cls.tmdbid == tmdbid,
|
||||
cls.type == mtype,
|
||||
cls.seasons == season,
|
||||
cls.episodes == episode,
|
||||
cls.dest == dest).all()
|
||||
# 电视剧某季
|
||||
elif season:
|
||||
elif season is not None:
|
||||
return db.query(cls).filter(cls.tmdbid == tmdbid,
|
||||
cls.type == mtype,
|
||||
cls.seasons == season).all()
|
||||
@@ -290,14 +290,14 @@ class TransferHistory(Base):
|
||||
# 标题 + 年份
|
||||
elif title and year:
|
||||
# 电视剧某季某集
|
||||
if season and episode:
|
||||
if season is not None and episode:
|
||||
return db.query(cls).filter(cls.title == title,
|
||||
cls.year == year,
|
||||
cls.seasons == season,
|
||||
cls.episodes == episode,
|
||||
cls.dest == dest).all()
|
||||
# 电视剧某季
|
||||
elif season:
|
||||
elif season is not None:
|
||||
return db.query(cls).filter(cls.title == title,
|
||||
cls.year == year,
|
||||
cls.seasons == season).all()
|
||||
@@ -312,7 +312,7 @@ class TransferHistory(Base):
|
||||
return db.query(cls).filter(cls.title == title,
|
||||
cls.year == year).all()
|
||||
# 类型 + 转移路径(emby webhook season无tmdbid场景)
|
||||
elif mtype and season and dest:
|
||||
elif mtype and season is not None and dest:
|
||||
# 电视剧某季
|
||||
return db.query(cls).filter(cls.type == mtype,
|
||||
cls.seasons == season,
|
||||
|
||||
@@ -71,6 +71,7 @@ class SubscribeOper(DbOper):
|
||||
"backdrop": mediainfo.get_backdrop_image(),
|
||||
"vote": mediainfo.vote_average,
|
||||
"description": mediainfo.overview,
|
||||
"search_imdbid": 1 if kwargs.get('search_imdbid') else 0,
|
||||
"date": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
||||
})
|
||||
if not subscribe:
|
||||
@@ -91,7 +92,7 @@ class SubscribeOper(DbOper):
|
||||
判断是否存在
|
||||
"""
|
||||
if tmdbid:
|
||||
if season:
|
||||
if season is not None:
|
||||
return True if Subscribe.exists(self._db, tmdbid=tmdbid, season=season) else False
|
||||
else:
|
||||
return True if Subscribe.exists(self._db, tmdbid=tmdbid) else False
|
||||
@@ -111,6 +112,20 @@ class SubscribeOper(DbOper):
|
||||
"""
|
||||
return await Subscribe.async_get(self._db, rid=sid)
|
||||
|
||||
def get_by(self, type: str, season: Optional[str] = None, tmdbid: Optional[int] = None,
|
||||
doubanid: Optional[str] = None, bangumiid: Optional[str] = None) -> Optional[Subscribe]:
|
||||
"""
|
||||
根据条件查询订阅
|
||||
"""
|
||||
return Subscribe.get_by(self._db, type, season, tmdbid, doubanid, bangumiid)
|
||||
|
||||
async def async_get_by(self, type: str, season: Optional[str] = None, tmdbid: Optional[int] = None,
|
||||
doubanid: Optional[str] = None, bangumiid: Optional[str] = None) -> Optional[Subscribe]:
|
||||
"""
|
||||
根据条件查询订阅
|
||||
"""
|
||||
return await Subscribe.async_get_by(self._db, type, season, tmdbid, doubanid, bangumiid)
|
||||
|
||||
def list(self, state: Optional[str] = None) -> List[Subscribe]:
|
||||
"""
|
||||
获取订阅列表
|
||||
@@ -180,7 +195,7 @@ class SubscribeOper(DbOper):
|
||||
判断是否存在订阅历史
|
||||
"""
|
||||
if tmdbid:
|
||||
if season:
|
||||
if season is not None:
|
||||
return True if SubscribeHistory.exists(self._db, tmdbid=tmdbid, season=season) else False
|
||||
else:
|
||||
return True if SubscribeHistory.exists(self._db, tmdbid=tmdbid) else False
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import threading
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from app.db import DbOper
|
||||
@@ -17,6 +19,8 @@ class SystemConfigOper(DbOper, metaclass=Singleton):
|
||||
"""
|
||||
super().__init__()
|
||||
self.__SYSTEMCONF = {}
|
||||
self._rlock = threading.RLock()
|
||||
self._alock = asyncio.Lock()
|
||||
for item in SystemConfig.list(self._db):
|
||||
self.__SYSTEMCONF[item.key] = item.value
|
||||
|
||||
@@ -29,23 +33,24 @@ class SystemConfigOper(DbOper, metaclass=Singleton):
|
||||
"""
|
||||
if isinstance(key, SystemConfigKey):
|
||||
key = key.value
|
||||
# 旧值
|
||||
old_value = self.__SYSTEMCONF.get(key)
|
||||
# 更新内存(deepcopy避免内存共享)
|
||||
self.__SYSTEMCONF[key] = copy.deepcopy(value)
|
||||
conf = SystemConfig.get_by_key(self._db, key)
|
||||
if conf:
|
||||
if old_value != value:
|
||||
if value:
|
||||
conf.update(self._db, {"value": value})
|
||||
else:
|
||||
conf.delete(self._db, conf.id)
|
||||
with self._rlock:
|
||||
# 旧值
|
||||
old_value = self.__SYSTEMCONF.get(key)
|
||||
# 更新内存(deepcopy避免内存共享)
|
||||
self.__SYSTEMCONF[key] = copy.deepcopy(value)
|
||||
conf = SystemConfig.get_by_key(self._db, key)
|
||||
if conf:
|
||||
if old_value != value:
|
||||
if value:
|
||||
conf.update(self._db, {"value": value})
|
||||
else:
|
||||
conf.delete(self._db, conf.id)
|
||||
return True
|
||||
return None
|
||||
else:
|
||||
conf = SystemConfig(key=key, value=value)
|
||||
conf.create(self._db)
|
||||
return True
|
||||
return None
|
||||
else:
|
||||
conf = SystemConfig(key=key, value=value)
|
||||
conf.create(self._db)
|
||||
return True
|
||||
|
||||
async def async_set(self, key: Union[str, SystemConfigKey], value: Any) -> Optional[bool]:
|
||||
"""
|
||||
@@ -56,22 +61,32 @@ class SystemConfigOper(DbOper, metaclass=Singleton):
|
||||
"""
|
||||
if isinstance(key, SystemConfigKey):
|
||||
key = key.value
|
||||
# 旧值
|
||||
old_value = self.__SYSTEMCONF.get(key)
|
||||
# 更新内存(deepcopy避免内存共享)
|
||||
self.__SYSTEMCONF[key] = copy.deepcopy(value)
|
||||
conf = await SystemConfig.async_get_by_key(self._db, key)
|
||||
if conf:
|
||||
if old_value != value:
|
||||
async with self._alock:
|
||||
conf = await SystemConfig.async_get_by_key(self._db, key)
|
||||
# 确定是否需要更新数据库
|
||||
needs_db_update = False
|
||||
if conf:
|
||||
if conf.value != value:
|
||||
needs_db_update = True
|
||||
else: # 记录不存在,总是需要创建/更新
|
||||
needs_db_update = True
|
||||
if not needs_db_update:
|
||||
# 即使数据库值相同,也要确保缓存同步
|
||||
with self._rlock:
|
||||
self.__SYSTEMCONF[key] = copy.deepcopy(value)
|
||||
return None
|
||||
# 执行数据库更新
|
||||
if conf:
|
||||
if value:
|
||||
conf.update(self._db, {"value": value})
|
||||
await conf.async_update(self._db, {"value": value})
|
||||
else:
|
||||
conf.delete(self._db, conf.id)
|
||||
return True
|
||||
return None
|
||||
else:
|
||||
conf = SystemConfig(key=key, value=value)
|
||||
await conf.async_create(self._db)
|
||||
await conf.async_delete(self._db, conf.id)
|
||||
else:
|
||||
conf = SystemConfig(key=key, value=value)
|
||||
await conf.async_create(self._db)
|
||||
# 数据库更新成功后,再更新缓存
|
||||
with self._rlock:
|
||||
self.__SYSTEMCONF[key] = copy.deepcopy(value)
|
||||
return True
|
||||
|
||||
def get(self, key: Union[str, SystemConfigKey] = None) -> Any:
|
||||
@@ -82,15 +97,17 @@ class SystemConfigOper(DbOper, metaclass=Singleton):
|
||||
key = key.value
|
||||
if not key:
|
||||
return self.all()
|
||||
# 避免将__SYSTEMCONF内的值引用出去,会导致set时误判没有变动
|
||||
return copy.deepcopy(self.__SYSTEMCONF.get(key))
|
||||
with self._rlock:
|
||||
# 避免将__SYSTEMCONF内的值引用出去,会导致set时误判没有变动
|
||||
return copy.deepcopy(self.__SYSTEMCONF.get(key))
|
||||
|
||||
def all(self):
|
||||
"""
|
||||
获取所有系统设置
|
||||
"""
|
||||
# 避免将__SYSTEMCONF内的值引用出去,会导致set时误判没有变动
|
||||
return copy.deepcopy(self.__SYSTEMCONF)
|
||||
with self._rlock:
|
||||
# 避免将__SYSTEMCONF内的值引用出去,会导致set时误判没有变动
|
||||
return copy.deepcopy(self.__SYSTEMCONF)
|
||||
|
||||
def delete(self, key: Union[str, SystemConfigKey]) -> bool:
|
||||
"""
|
||||
@@ -98,10 +115,11 @@ class SystemConfigOper(DbOper, metaclass=Singleton):
|
||||
"""
|
||||
if isinstance(key, SystemConfigKey):
|
||||
key = key.value
|
||||
# 更新内存
|
||||
self.__SYSTEMCONF.pop(key, None)
|
||||
# 写入数据库
|
||||
conf = SystemConfig.get_by_key(self._db, key)
|
||||
if conf:
|
||||
conf.delete(self._db, conf.id)
|
||||
return True
|
||||
with self._rlock:
|
||||
# 更新内存
|
||||
self.__SYSTEMCONF.pop(key, None)
|
||||
# 写入数据库
|
||||
conf = SystemConfig.get_by_key(self._db, key)
|
||||
if conf:
|
||||
conf.delete(self._db, conf.id)
|
||||
return True
|
||||
|
||||
@@ -125,7 +125,7 @@ class TransferHistoryOper(DbOper):
|
||||
"""
|
||||
新增转移成功历史记录
|
||||
"""
|
||||
self.add_force(
|
||||
return self.add_force(
|
||||
src=fileitem.path,
|
||||
src_storage=fileitem.storage,
|
||||
src_fileitem=fileitem.model_dump(),
|
||||
|
||||
@@ -19,41 +19,42 @@ class CookieHelper:
|
||||
"username": [
|
||||
'//input[@name="username"]',
|
||||
'//input[@id="form_item_username"]',
|
||||
'//input[@id="username"]'
|
||||
'//input[@id="username"]',
|
||||
],
|
||||
"password": [
|
||||
'//input[@name="password"]',
|
||||
'//input[@id="form_item_password"]',
|
||||
'//input[@id="password"]',
|
||||
'//input[@type="password"]'
|
||||
'//input[@type="password"]',
|
||||
],
|
||||
"captcha": [
|
||||
'//input[@name="imagestring"]',
|
||||
'//input[@name="captcha"]',
|
||||
'//input[@id="form_item_captcha"]',
|
||||
'//input[@placeholder="驗證碼"]'
|
||||
'//input[@placeholder="驗證碼"]',
|
||||
],
|
||||
"captcha_img": [
|
||||
'//img[@alt="captcha"]/@src',
|
||||
'//img[@alt="CAPTCHA"]/@src',
|
||||
'//img[@alt="SECURITY CODE"]/@src',
|
||||
'//img[@id="LAY-user-get-vercode"]/@src',
|
||||
'//img[contains(@src,"/api/getCaptcha")]/@src'
|
||||
'//img[contains(@src,"/api/getCaptcha")]/@src',
|
||||
],
|
||||
"submit": [
|
||||
'//input[@type="submit"]',
|
||||
'//button[@type="submit"]',
|
||||
'//button[@lay-filter="login"]',
|
||||
'//button[@lay-filter="formLogin"]',
|
||||
'//input[@type="button"][@value="登录"]'
|
||||
'//input[@type="button"][@value="登录"]',
|
||||
'//input[@id="submit-btn"]',
|
||||
],
|
||||
"error": [
|
||||
"//table[@class='main']//td[@class='text']/text()"
|
||||
"//table[@class='main']//td[@class='text']/text()",
|
||||
],
|
||||
"twostep": [
|
||||
'//input[@name="two_step_code"]',
|
||||
'//input[@name="2fa_secret"]',
|
||||
'//input[@name="otp"]'
|
||||
'//input[@name="otp"]',
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@@ -142,19 +142,23 @@ class DirectoryHelper:
|
||||
# 计算重命名中的文件夹层数
|
||||
rename_list = rename_format.split("/")
|
||||
rename_format_level = len(rename_list) - 1
|
||||
# 查找标题参数所在层
|
||||
for level, name in enumerate(rename_list):
|
||||
# 反向查找标题参数所在层
|
||||
for level, name in enumerate(reversed(rename_list)):
|
||||
if level == 0:
|
||||
# 跳过文件名的标题参数
|
||||
continue
|
||||
matchs = JINJA2_VAR_PATTERN.findall(name)
|
||||
if not matchs:
|
||||
continue
|
||||
# 处理特例,有的人重命名的第一层是年份、分辨率
|
||||
if any("title" in m for m in matchs):
|
||||
# 找出含标题的这一层作为媒体根目录
|
||||
rename_format_level -= level
|
||||
if (any("title" in m for m in matchs)
|
||||
and not any("season" in m for m in matchs)):
|
||||
# 找出最后一层含有标题且不含季参数的目录作为媒体根目录
|
||||
rename_format_level = level
|
||||
break
|
||||
else:
|
||||
# 假定第一层目录是媒体根目录
|
||||
logger.warn(f"重命名格式 {rename_format} 缺少标题参数")
|
||||
logger.warn(f"重命名格式 {rename_format} 缺少标题目录")
|
||||
if rename_format_level > len(rename_path.parents):
|
||||
# 通常因为路径以/结尾,被Path规范化删除了
|
||||
logger.error(f"路径 {rename_path} 不匹配重命名格式 {rename_format}")
|
||||
|
||||
@@ -25,7 +25,7 @@ class DownloaderHelper(ServiceBaseHelper[DownloaderConf]):
|
||||
) -> bool:
|
||||
"""
|
||||
通用的下载器类型判断方法
|
||||
:param service_type: 下载器的类型名称(如 'qbittorrent', 'transmission')
|
||||
:param service_type: 下载器的类型名称(如 'qbittorrent', 'transmission', 'rtorrent')
|
||||
:param service: 要判断的服务信息
|
||||
:param name: 服务的名称
|
||||
:return: 如果服务类型或实例为指定类型,返回 True;否则返回 False
|
||||
|
||||
@@ -1,13 +1,91 @@
|
||||
"""LLM模型相关辅助功能"""
|
||||
|
||||
from typing import List
|
||||
|
||||
from app.core.config import settings
|
||||
from app.log import logger
|
||||
|
||||
|
||||
class LLMHelper:
|
||||
"""LLM模型相关辅助功能"""
|
||||
|
||||
def get_models(self, provider: str, api_key: str, base_url: str = None) -> List[str]:
|
||||
@staticmethod
|
||||
def get_llm(streaming: bool = False):
|
||||
"""
|
||||
获取LLM实例
|
||||
:param streaming: 是否启用流式输出
|
||||
:return: LLM实例
|
||||
"""
|
||||
provider = settings.LLM_PROVIDER.lower()
|
||||
api_key = settings.LLM_API_KEY
|
||||
|
||||
if not api_key:
|
||||
raise ValueError("未配置LLM API Key")
|
||||
|
||||
if provider == "google":
|
||||
if settings.PROXY_HOST:
|
||||
# 通过代理使用 Google 的 OpenAI 兼容接口
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
model = ChatOpenAI(
|
||||
model=settings.LLM_MODEL,
|
||||
api_key=api_key,
|
||||
max_retries=3,
|
||||
base_url="https://generativelanguage.googleapis.com/v1beta/openai",
|
||||
temperature=settings.LLM_TEMPERATURE,
|
||||
streaming=streaming,
|
||||
stream_usage=True,
|
||||
openai_proxy=settings.PROXY_HOST,
|
||||
)
|
||||
else:
|
||||
# 使用 langchain-google-genai 原生接口(v4 API 变更:google_api_key → api_key,max_retries → retries)
|
||||
from langchain_google_genai import ChatGoogleGenerativeAI
|
||||
|
||||
model = ChatGoogleGenerativeAI(
|
||||
model=settings.LLM_MODEL,
|
||||
api_key=api_key,
|
||||
retries=3,
|
||||
temperature=settings.LLM_TEMPERATURE,
|
||||
streaming=streaming
|
||||
)
|
||||
elif provider == "deepseek":
|
||||
from langchain_deepseek import ChatDeepSeek
|
||||
|
||||
model = ChatDeepSeek(
|
||||
model=settings.LLM_MODEL,
|
||||
api_key=api_key,
|
||||
max_retries=3,
|
||||
temperature=settings.LLM_TEMPERATURE,
|
||||
streaming=streaming,
|
||||
stream_usage=True,
|
||||
)
|
||||
else:
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
model = ChatOpenAI(
|
||||
model=settings.LLM_MODEL,
|
||||
api_key=api_key,
|
||||
max_retries=3,
|
||||
base_url=settings.LLM_BASE_URL,
|
||||
temperature=settings.LLM_TEMPERATURE,
|
||||
streaming=streaming,
|
||||
stream_usage=True,
|
||||
openai_proxy=settings.PROXY_HOST,
|
||||
)
|
||||
|
||||
# 检查是否有profile
|
||||
if hasattr(model, "profile") and model.profile:
|
||||
logger.info(f"使用LLM模型: {model.model},Profile: {model.profile}")
|
||||
else:
|
||||
model.profile = {
|
||||
"max_input_tokens": settings.LLM_MAX_CONTEXT_TOKENS * 1000, # 转换为token单位
|
||||
}
|
||||
|
||||
return model
|
||||
|
||||
def get_models(
|
||||
self, provider: str, api_key: str, base_url: str = None
|
||||
) -> List[str]:
|
||||
"""获取模型列表"""
|
||||
logger.info(f"获取 {provider} 模型列表...")
|
||||
if provider == "google":
|
||||
@@ -17,18 +95,25 @@ class LLMHelper:
|
||||
|
||||
@staticmethod
|
||||
def _get_google_models(api_key: str) -> List[str]:
|
||||
"""获取Google模型列表"""
|
||||
"""获取Google模型列表(使用 google-genai SDK v1)"""
|
||||
try:
|
||||
import google.generativeai as genai
|
||||
genai.configure(api_key=api_key)
|
||||
models = genai.list_models()
|
||||
return [m.name for m in models if 'generateContent' in m.supported_generation_methods]
|
||||
from google import genai
|
||||
|
||||
client = genai.Client(api_key=api_key)
|
||||
models = client.models.list()
|
||||
return [
|
||||
m.name
|
||||
for m in models
|
||||
if m.supported_actions and "generateContent" in m.supported_actions
|
||||
]
|
||||
except Exception as e:
|
||||
logger.error(f"获取Google模型列表失败:{e}")
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def _get_openai_compatible_models(provider: str, api_key: str, base_url: str = None) -> List[str]:
|
||||
def _get_openai_compatible_models(
|
||||
provider: str, api_key: str, base_url: str = None
|
||||
) -> List[str]:
|
||||
"""获取OpenAI兼容模型列表"""
|
||||
try:
|
||||
from openai import OpenAI
|
||||
|
||||
@@ -164,6 +164,8 @@ class TemplateContextBuilder:
|
||||
"part": meta.part,
|
||||
# 自定义占位符
|
||||
"customization": meta.customization,
|
||||
# fps
|
||||
"fps": meta.fps,
|
||||
}
|
||||
|
||||
tech_metadata = {
|
||||
@@ -539,7 +541,7 @@ class MessageTemplateHelper:
|
||||
获取消息模板
|
||||
"""
|
||||
template_dict: dict[str, str] = SystemConfigOper().get(SystemConfigKey.NotificationTemplates)
|
||||
return template_dict.get(f"{message.ctype.value}")
|
||||
return template_dict.get(message.ctype.value)
|
||||
|
||||
|
||||
class MessageQueueManager(metaclass=SingletonClass):
|
||||
|
||||
@@ -90,6 +90,79 @@ class PassKeyHelper:
|
||||
logger.error(f"标准化凭证ID失败: {e}")
|
||||
return credential_id
|
||||
|
||||
@staticmethod
|
||||
def _base64_encode_urlsafe(data: bytes) -> str:
|
||||
"""
|
||||
Base64 URL Safe 编码(不带填充)
|
||||
|
||||
:param data: 要编码的字节数据
|
||||
:return: Base64 URL Safe 编码的字符串
|
||||
"""
|
||||
return base64.urlsafe_b64encode(data).decode('utf-8').rstrip('=')
|
||||
|
||||
@staticmethod
|
||||
def _base64_decode_urlsafe(data: str) -> bytes:
|
||||
"""
|
||||
Base64 URL Safe 解码(自动添加填充)
|
||||
|
||||
:param data: Base64 URL Safe 编码的字符串
|
||||
:return: 解码后的字节数据
|
||||
"""
|
||||
return base64.urlsafe_b64decode(data + '==')
|
||||
|
||||
@staticmethod
|
||||
def _parse_credential_list(credentials: List[Dict[str, Any]]) -> List[PublicKeyCredentialDescriptor]:
|
||||
"""
|
||||
解析凭证列表为 PublicKeyCredentialDescriptor 列表
|
||||
|
||||
:param credentials: 凭证字典列表
|
||||
:return: PublicKeyCredentialDescriptor 列表
|
||||
"""
|
||||
result = []
|
||||
for cred in credentials:
|
||||
try:
|
||||
result.append(
|
||||
PublicKeyCredentialDescriptor(
|
||||
id=PassKeyHelper._base64_decode_urlsafe(cred['credential_id']),
|
||||
transports=[
|
||||
AuthenticatorTransport(t) for t in cred.get('transports', '').split(',') if t
|
||||
] if cred.get('transports') else None
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"解析凭证失败: {e}")
|
||||
continue
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _get_user_verification_requirement(user_verification: Optional[str] = None) -> UserVerificationRequirement:
|
||||
"""
|
||||
获取用户验证要求
|
||||
|
||||
:param user_verification: 指定的用户验证要求,如果不指定则从配置中读取
|
||||
:return: UserVerificationRequirement
|
||||
"""
|
||||
if user_verification:
|
||||
return UserVerificationRequirement(user_verification)
|
||||
return UserVerificationRequirement.REQUIRED if settings.PASSKEY_REQUIRE_UV \
|
||||
else UserVerificationRequirement.PREFERRED
|
||||
|
||||
@staticmethod
|
||||
def _get_verification_params(
|
||||
expected_origin: Optional[str] = None,
|
||||
expected_rp_id: Optional[str] = None
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
获取验证参数(origin 和 rp_id)
|
||||
|
||||
:param expected_origin: 期望的源地址
|
||||
:param expected_rp_id: 期望的RP ID
|
||||
:return: (origin, rp_id)
|
||||
"""
|
||||
origin = expected_origin or PassKeyHelper.get_origin()
|
||||
rp_id = expected_rp_id or PassKeyHelper.get_rp_id()
|
||||
return origin, rp_id
|
||||
|
||||
@staticmethod
|
||||
def generate_registration_options(
|
||||
user_id: int,
|
||||
@@ -109,27 +182,13 @@ class PassKeyHelper:
|
||||
try:
|
||||
# 用户信息
|
||||
user_id_bytes = str(user_id).encode('utf-8')
|
||||
|
||||
|
||||
# 排除已有的凭证
|
||||
exclude_credentials = []
|
||||
if existing_credentials:
|
||||
for cred in existing_credentials:
|
||||
try:
|
||||
exclude_credentials.append(
|
||||
PublicKeyCredentialDescriptor(
|
||||
id=base64.urlsafe_b64decode(cred['credential_id'] + '=='),
|
||||
transports=[
|
||||
AuthenticatorTransport(t) for t in cred.get('transports', '').split(',') if t
|
||||
] if cred.get('transports') else None
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"解析凭证失败: {e}")
|
||||
continue
|
||||
exclude_credentials = PassKeyHelper._parse_credential_list(existing_credentials) \
|
||||
if existing_credentials else None
|
||||
|
||||
# 用户验证要求
|
||||
uv_requirement = UserVerificationRequirement.REQUIRED if settings.PASSKEY_REQUIRE_UV \
|
||||
else UserVerificationRequirement.PREFERRED
|
||||
uv_requirement = PassKeyHelper._get_user_verification_requirement()
|
||||
|
||||
# 生成注册选项
|
||||
options = generate_registration_options(
|
||||
@@ -138,7 +197,7 @@ class PassKeyHelper:
|
||||
user_id=user_id_bytes,
|
||||
user_name=username,
|
||||
user_display_name=display_name or username,
|
||||
exclude_credentials=exclude_credentials if exclude_credentials else None,
|
||||
exclude_credentials=exclude_credentials,
|
||||
authenticator_selection=AuthenticatorSelectionCriteria(
|
||||
authenticator_attachment=None,
|
||||
resident_key=ResidentKeyRequirement.REQUIRED,
|
||||
@@ -152,9 +211,9 @@ class PassKeyHelper:
|
||||
|
||||
# 转换为JSON
|
||||
options_json = options_to_json(options)
|
||||
|
||||
|
||||
# 提取challenge(用于后续验证)
|
||||
challenge = base64.urlsafe_b64encode(options.challenge).decode('utf-8').rstrip('=')
|
||||
challenge = PassKeyHelper._base64_encode_urlsafe(options.challenge)
|
||||
|
||||
return options_json, challenge
|
||||
|
||||
@@ -162,29 +221,6 @@ class PassKeyHelper:
|
||||
logger.error(f"生成注册选项失败: {e}")
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _get_verified_origin(credential: Dict[str, Any], rp_id: str, default_origin: str) -> str:
|
||||
"""
|
||||
在 localhost 环境下获取并验证实际 Origin,否则返回默认值
|
||||
"""
|
||||
if not settings.APP_DOMAIN and rp_id == 'localhost':
|
||||
try:
|
||||
# 解析 clientDataJSON 获取实际的 origin
|
||||
client_data_json = json.loads(
|
||||
base64.urlsafe_b64decode(
|
||||
credential['response']['clientDataJSON'].replace('-', '+').replace('_', '/') + '=='
|
||||
).decode('utf-8')
|
||||
)
|
||||
actual_origin = client_data_json.get('origin', '')
|
||||
hostname = urlparse(actual_origin).hostname
|
||||
|
||||
if hostname in ['localhost', '127.0.0.1']:
|
||||
logger.info(f"本地环境,使用动态 origin: {actual_origin}")
|
||||
return actual_origin
|
||||
except Exception as e:
|
||||
logger.warning(f"无法提取动态 origin: {e}")
|
||||
return default_origin
|
||||
|
||||
@staticmethod
|
||||
def verify_registration_response(
|
||||
credential: Dict[str, Any],
|
||||
@@ -203,18 +239,13 @@ class PassKeyHelper:
|
||||
"""
|
||||
try:
|
||||
# 准备验证参数
|
||||
origin = expected_origin or PassKeyHelper.get_origin()
|
||||
rp_id = expected_rp_id or PassKeyHelper.get_rp_id()
|
||||
|
||||
origin, rp_id = PassKeyHelper._get_verification_params(expected_origin, expected_rp_id)
|
||||
# 解码challenge
|
||||
challenge_bytes = base64.urlsafe_b64decode(expected_challenge + '==')
|
||||
challenge_bytes = PassKeyHelper._base64_decode_urlsafe(expected_challenge)
|
||||
|
||||
# 构建RegistrationCredential对象
|
||||
registration_credential = parse_registration_credential_json(json.dumps(credential))
|
||||
|
||||
# 获取并验证 Origin
|
||||
origin = PassKeyHelper._get_verified_origin(credential, rp_id, origin)
|
||||
|
||||
# 验证注册响应
|
||||
verification = verify_registration_response(
|
||||
credential=registration_credential,
|
||||
@@ -225,8 +256,8 @@ class PassKeyHelper:
|
||||
)
|
||||
|
||||
# 提取信息
|
||||
credential_id = base64.urlsafe_b64encode(verification.credential_id).decode('utf-8').rstrip('=')
|
||||
public_key = base64.urlsafe_b64encode(verification.credential_public_key).decode('utf-8').rstrip('=')
|
||||
credential_id = PassKeyHelper._base64_encode_urlsafe(verification.credential_id)
|
||||
public_key = PassKeyHelper._base64_encode_urlsafe(verification.credential_public_key)
|
||||
sign_count = verification.sign_count
|
||||
# aaguid 可能已经是字符串格式,也可能是bytes
|
||||
if verification.aaguid:
|
||||
@@ -257,41 +288,24 @@ class PassKeyHelper:
|
||||
"""
|
||||
try:
|
||||
# 允许的凭证
|
||||
allow_credentials = []
|
||||
if existing_credentials:
|
||||
for cred in existing_credentials:
|
||||
try:
|
||||
allow_credentials.append(
|
||||
PublicKeyCredentialDescriptor(
|
||||
id=base64.urlsafe_b64decode(cred['credential_id'] + '=='),
|
||||
transports=[
|
||||
AuthenticatorTransport(t) for t in cred.get('transports', '').split(',') if t
|
||||
] if cred.get('transports') else None
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"解析凭证失败: {e}")
|
||||
continue
|
||||
allow_credentials = PassKeyHelper._parse_credential_list(existing_credentials) \
|
||||
if existing_credentials else None
|
||||
|
||||
# 用户验证要求
|
||||
if not user_verification:
|
||||
uv_requirement = UserVerificationRequirement.REQUIRED if settings.PASSKEY_REQUIRE_UV \
|
||||
else UserVerificationRequirement.PREFERRED
|
||||
else:
|
||||
uv_requirement = UserVerificationRequirement(user_verification)
|
||||
uv_requirement = PassKeyHelper._get_user_verification_requirement(user_verification)
|
||||
|
||||
# 生成认证选项
|
||||
options = generate_authentication_options(
|
||||
rp_id=PassKeyHelper.get_rp_id(),
|
||||
allow_credentials=allow_credentials if allow_credentials else None,
|
||||
allow_credentials=allow_credentials,
|
||||
user_verification=uv_requirement
|
||||
)
|
||||
|
||||
# 转换为JSON
|
||||
options_json = options_to_json(options)
|
||||
|
||||
|
||||
# 提取challenge
|
||||
challenge = base64.urlsafe_b64encode(options.challenge).decode('utf-8').rstrip('=')
|
||||
challenge = PassKeyHelper._base64_encode_urlsafe(options.challenge)
|
||||
|
||||
return options_json, challenge
|
||||
|
||||
@@ -321,19 +335,14 @@ class PassKeyHelper:
|
||||
"""
|
||||
try:
|
||||
# 准备验证参数
|
||||
origin = expected_origin or PassKeyHelper.get_origin()
|
||||
rp_id = expected_rp_id or PassKeyHelper.get_rp_id()
|
||||
|
||||
origin, rp_id = PassKeyHelper._get_verification_params(expected_origin, expected_rp_id)
|
||||
# 解码
|
||||
challenge_bytes = base64.urlsafe_b64decode(expected_challenge + '==')
|
||||
public_key_bytes = base64.urlsafe_b64decode(credential_public_key + '==')
|
||||
challenge_bytes = PassKeyHelper._base64_decode_urlsafe(expected_challenge)
|
||||
public_key_bytes = PassKeyHelper._base64_decode_urlsafe(credential_public_key)
|
||||
|
||||
# 构建AuthenticationCredential对象
|
||||
authentication_credential = parse_authentication_credential_json(json.dumps(credential))
|
||||
|
||||
# 获取并验证 Origin
|
||||
origin = PassKeyHelper._get_verified_origin(credential, rp_id, origin)
|
||||
|
||||
# 验证认证响应
|
||||
verification = verify_authentication_response(
|
||||
credential=authentication_credential,
|
||||
|
||||
@@ -13,9 +13,10 @@ import aiofiles
|
||||
import aioshutil
|
||||
import httpx
|
||||
from anyio import Path as AsyncPath
|
||||
from packaging.requirements import Requirement
|
||||
from packaging.specifiers import SpecifierSet, InvalidSpecifier
|
||||
from packaging.version import Version, InvalidVersion
|
||||
from pkg_resources import Requirement, working_set
|
||||
from importlib.metadata import distributions
|
||||
from requests import Response
|
||||
|
||||
from app.core.cache import cached
|
||||
@@ -729,18 +730,26 @@ class PluginHelper(metaclass=WeakSingleton):
|
||||
def __get_installed_packages(self) -> Dict[str, Version]:
|
||||
"""
|
||||
获取已安装的包及其版本
|
||||
使用 pkg_resources 获取当前环境中已安装的包,标准化包名并转换版本信息
|
||||
使用 importlib.metadata 获取当前环境中已安装的包,标准化包名并转换版本信息
|
||||
对于无法解析的版本,记录警告日志并跳过
|
||||
:return: 已安装包的字典,格式为 {package_name: Version}
|
||||
"""
|
||||
installed_packages = {}
|
||||
try:
|
||||
for dist in working_set:
|
||||
pkg_name = self.__standardize_pkg_name(dist.project_name)
|
||||
for dist in distributions():
|
||||
name = dist.metadata.get("Name")
|
||||
if not name:
|
||||
continue
|
||||
pkg_name = self.__standardize_pkg_name(name)
|
||||
version_str = dist.metadata.get("Version") or getattr(dist, "version", None)
|
||||
if not version_str:
|
||||
continue
|
||||
try:
|
||||
installed_packages[pkg_name] = Version(dist.version)
|
||||
v = Version(version_str)
|
||||
if pkg_name not in installed_packages or v > installed_packages[pkg_name]:
|
||||
installed_packages[pkg_name] = v
|
||||
except InvalidVersion:
|
||||
logger.debug(f"无法解析已安装包 '{pkg_name}' 的版本:{dist.version}")
|
||||
logger.debug(f"无法解析已安装包 '{pkg_name}' 的版本:{version_str}")
|
||||
continue
|
||||
return installed_packages
|
||||
except Exception as e:
|
||||
@@ -844,12 +853,14 @@ class PluginHelper(metaclass=WeakSingleton):
|
||||
@staticmethod
|
||||
def __standardize_pkg_name(name: str) -> str:
|
||||
"""
|
||||
标准化包名,将包名转换为小写并将连字符替换为下划线
|
||||
标准化包名,将包名转换为小写,连字符与点替换为下划线(与 PEP 503 归一化风格一致)
|
||||
|
||||
:param name: 原始包名
|
||||
:return: 标准化后的包名
|
||||
"""
|
||||
return name.lower().replace("-", "_") if name else name
|
||||
if not name:
|
||||
return name
|
||||
return name.lower().replace("-", "_").replace(".", "_")
|
||||
|
||||
async def async_get_plugin_package_version(self, pid: str, repo_url: str,
|
||||
package_version: Optional[str] = None) -> Optional[str]:
|
||||
|
||||
@@ -3,10 +3,9 @@ from typing import Union, Optional
|
||||
|
||||
from app.core.cache import TTLCache
|
||||
from app.schemas.types import ProgressKey
|
||||
from app.utils.singleton import WeakSingleton
|
||||
|
||||
|
||||
class ProgressHelper(metaclass=WeakSingleton):
|
||||
class ProgressHelper:
|
||||
"""
|
||||
处理进度辅助类
|
||||
"""
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user